מפרט StableHLO

‫StableHLO הוא קבוצת פעולות לפעולות ברמה גבוהה (HLO) במודלים של למידת מכונה (ML). ‫StableHLO פועל כשכבת ניוד בין מסגרות שונות של למידת מכונה (ML) ומהדרים של למידת מכונה: מסגרות של למידת מכונה שמפיקות תוכניות StableHLO תואמות למהדרים של למידת מכונה שצורכים תוכניות StableHLO.

המטרה שלנו היא לפשט ולהאיץ את פיתוח ה-ML על ידי יצירת יכולת פעולה הדדית רבה יותר בין מסגרות שונות של ML (כמו TensorFlow, ‏ JAX ו-PyTorch) ובין קומפיילרים של ML (כמו XLA ו-IREE). לכן, במסמך הזה מפורטות ההגדרות של שפת התכנות StableHLO.

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

כדי לראות את המפרט מגרסה קודמת של StableHLO, פותחים את המאגר בגרסה המתויגת הרלוונטית. לדוגמה, StableHLO v0.19.0 Spec. כדי לראות את השינויים שחלו בכל עדכון של גרסה משנית של StableHLO, אפשר לעיין ביומן הגרסאות ב-VhloDialect.td.

תוכניות

Program ::= {Func}

תוכניות StableHLO מורכבות ממספר כלשהו של פונקציות StableHLO. למטה מופיעה דוגמה לתוכנית עם פונקציה @main שיש לה 3 קלטים (%image, ‏ %weights ו-%bias) ופלט אחד. הגוף של הפונקציה כולל 6 פעולות.

func.func @main(
  %image: tensor<28x28xf32>,
  %weights: tensor<784x10xf32>,
  %bias: tensor<1x10xf32>
) -> tensor<1x10xf32> {
  %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
  %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
  %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
  %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  "func.return"(%4): (tensor<1x10xf32>) -> ()
}

פונקציות

Func        ::= 'func' '.' 'func' FuncId FuncInputs FuncOutputs '{' FuncBody '}'
FuncInputs  ::= '(' [FuncInput {',' FuncInput}] `)`
FuncInput   ::= ValueId ':' ValueType
FuncOutputs ::= ['->' FuncOutput, {',' FuncOutput}]
FuncOutput  ::= ValueType
FuncBody    ::= {Op}

פונקציות StableHLO (שנקראות גם פונקציות בעלות שם) כוללות מזהה, קלט/פלט וגוף. בעתיד, אנחנו מתכננים להוסיף מטא-נתונים לפונקציות כדי לשפר את התאימות ל-HLO (‎#425,‏ ‎#626,‏ ‎#740,‏ ‎#744).

מזהים

FuncId  ::= '@' letter {letter | digit}
ValueId ::= '%' digit {digit}
          | '%' letter {letter | digit}
letter  ::= 'a' | ... | 'z' | 'A' | ... | 'Z' | '_'
digit   ::= '0' | ... | '9'

מזהי StableHLO דומים למזהים בהרבה שפות תכנות, עם שני מאפיינים ייחודיים: 1) לכל המזהים יש סיגילים שמבחינים בין סוגים שונים של מזהים, 2) מזהי ערכים יכולים להיות מספריים לחלוטין כדי לפשט את יצירת התוכניות של StableHLO.

סוגים

Type         ::= ValueType | NonValueType
ValueType    ::= TensorType | QuantizedTensorType | TokenType | TupleType | BufferType
NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType

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

TensorType ::= 'tensor' '<' Shape TensorElementType '>'
Shape ::= {DimensionSize 'x'}
DimensionSize ::= digit {digit} | '?'

סוגי טנסורים מייצגים טנסורים, כלומר מערכים רב-ממדיים. יש להם צורה וסוג רכיב, כאשר צורה מייצגת גדלים של ממדים לא שליליים או לא ידועים בסדר עולה של הממדים המתאימים (שנקראים גם צירים) שממוספרים מ-0 עד R-1. מספר המאפיינים R נקרא דרגה. לדוגמה, tensor<2x3xf32> הוא סוג טנסור עם צורה 2x3 וסוג אלמנט f32. יש לו שני מאפיינים (או במילים אחרות, שני צירים) – מאפיין 0 ומאפיין 1 – שהגדלים שלהם הם 2 ו-3. הדירוג שלו הוא 2.

יכול להיות שצורות יהיו לא ידועות באופן חלקי או מלא (דינמיות), למשל tensor<?x2xf64> לא ידוע באופן חלקי ו-tensor<?x?xf64> לא ידוע באופן מלא. גדלים של מאפיינים דינמיים מיוצגים באמצעות ?. אי אפשר לבטל את הדירוג של צורות.

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

QuantizedTensorType ::= 'tensor' '<' Shape QuantizedTensorElementType '>'
QuantizedTensorElementType ::= '!quant.uniform' '<'
                  QuantizationStorageType
                  ['<' QuantizationStorageMin ':' QuantizationStorageMax '>']
                  ':' QuantizationExpressedType
                  [':' QuantizationDimension]
                  ',' QuantizationParameters '>'
QuantizationStorageType ::= IntegerType
QuantizationStorageMin ::= IntegerLiteral
QuantizationStorageMax ::= IntegerLiteral
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerLiteral
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale [':' QuantizationZeroPoint]
QuantizationScale ::= FloatLiteral
QuantizationZeroPoint ::= IntegerLiteral

סוגי רכיבים שעברו קוונטיזציה מייצגים ערכים שלמים של סוג אחסון בטווח storage_min עד storage_max (כולל) שתואמים לערכים של נקודה צפה של סוג מבוטא. עבור ערך שלם נתון i, אפשר לחשב את הערך התואם בנקודה צפה f באופן הבא: f = (i - zero_point) * scale, כאשר scale ו-zero_point נקראים פרמטרים של קוונטיזציה. המאפיינים storage_min ו-storage_max הם אופציונליים בתחביר, אבל ערכי ברירת המחדל שלהם הם min_value(storage_type) ו-max_value(storage_type) בהתאמה. יש מגבלות על סוגי רכיבים שעברו קוונטיזציה:

• (C1) type(storage_min) = storage_type.
• ‫(C2) type(storage_max) = storage_type.
• ‫(C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type).
• ‫(C4) type(scales...) = expressed_type.
• ‫(C5) 0 < scales.
• ‫(C6) is_finite(scales...).
• ‫(C7) storage_min <= zero_points <= storage_max.
• ‫(C8) type(zero_points...) = storage_type.
• ‫(C9) size(scales) = size(zero_points).
• ‫(C10) אם is_empty(quantization_dimension), אז size(scales) = 1.
• ‫(C11) 0 <= quantization_dimension.

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

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

דיון נוסף מתנהל לגבי הסמנטיקה של QuantizationStorageMin ושל QuantizationStorageMax כדי לקבוע אם צריך להגביל את הערכים האלה ואת הערכים של טנסורים שעברו קוונטיזציה (‎#1406).

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

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

בטנסורים שעברו קוונטיזציה, הקוונטיזציה יכולה להיות per-tensor, כלומר, עם scale ו-zero_point אחדים לכל הטנסור, או per-axis, כלומר, עם כמה scales ו-zero_points, זוג אחד לכל פרוסה של מימד מסוים quantization_dimension. באופן רשמי יותר, בטנזור t עם כימות לכל ציר, יש dim(t, quantization_dimension) פרוסות של quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] וכו'. כל הרכיבים בפרוסה ה-i משתמשים ב-scales[i] וב-zero_points[i] כפרמטרים של הכימות. יש מגבלות על סוגי טנסורים שעברו קוונטיזציה:

• לכימות לכל טנסור:
  • אין אילוצים נוספים.
• לכימות לכל ציר:
  • (C12) quantization_dimension < rank(self).
  • ‫(C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

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

TupleType ::= 'tuple' '<' TupleElementTypes '>'
TupleElementTypes ::= [ValueType {',' ValueType}]

סוגי מאגרים מייצגים מאגרים. לדוגמה, ב-XLA, מאגרי נתונים זמניים הם מערכים רב-ממדיים עם אחסון עקבי. בדומה לסוגי טנסורים, לסוגי מאגרים יש צורה וסוג רכיב, כאשר הצורה מייצגת גדלים של ממדים לא שליליים או לא ידועים בסדר עולה של הממדים המתאימים (שנקראים גם צירים) שממוספרים מ-0 עד R-1. מספר המאפיינים R נקרא דרגה. לדוגמה, memref<2x3xf32> הוא סוג של מאגר עם צורה 2x3 וסוג רכיב f32. יש לו שני ממדים (או במילים אחרות, שני צירים) – ממד 0 וממד 1 – שהגדלים שלהם הם 2 ו-3. הדירוג שלו הוא 2.

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

סוגי טאפלים מייצגים טאפלים, כלומר רשימות הטרוגניות. הטופלים הם תכונה מדור קודם שקיימת רק לצורך תאימות ל-HLO. ב-HLO, נעשה שימוש בטפילים כדי לייצג קלטים ופלט משתנים. ב-StableHLO יש תמיכה מובנית בקלט ובפלט עם מספר משתנה של ארגומנטים, והשימוש היחיד בטפלים ב-StableHLO הוא לייצוג מקיף של HLO ABI, שבו למשל T, ‏ tuple<T> ו-tuple<tuple<T>> עשויים להיות שונים באופן משמעותי בהתאם להטמעה מסוימת. בעתיד, אנחנו מתכננים לבצע שינויים ב-HLO ABI, שיאפשרו לנו להסיר סוגי טאפל מ-StableHLO (‎#598).

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si2' | 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui2' | 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

סוגי רכיבים מייצגים רכיבים של סוגי טנסורים. בניגוד לשפות תכנות רבות, הסוגים האלה לא מוגדרים כסוגים ממחלקה ראשונה ב-StableHLO. המשמעות היא שתוכניות StableHLO לא יכולות לייצג ישירות ערכים מהסוגים האלה (כתוצאה מכך, מקובל לייצג ערכים סקלריים מהסוג T באמצעות ערכים טנסוריים אפסיים מהסוג tensor<T>).

• סוג בוליאני מייצג ערכים בוליאניים true ו-false.
• סוגי מספרים שלמים יכולים להיות עם סימן (si) או ללא סימן (ui), ויכולים להיות בעלי אחת מרוחבי הסיביות הנתמכים (2,‏ 4,‏ 8,‏ 16,‏ 32 או 64). סוגי siN עם סימן מייצגים ערכים של מספרים שלמים מ--2^(N-1) עד 2^(N-1)-1 כולל, וסוגי uiN ללא סימן מייצגים ערכים של מספרים שלמים מ-0 עד 2^N-1 כולל.
• סוגי נקודה צפה יכולים להיות אחד מהסוגים הבאים:
  • הערכים f8E3M4,‏ f8E4M3 ו-f8E5M2 הם מספרים שלמים של 8 ביט בפורמט נקודה צפה, בהתאם למוסכמות IEEE-754.
  • ‫f8E4M3FN ו-f8E5M2, בהתאמה לקידודים E4M3 ו-E5M2 של פורמט FP8 שמתואר במאמר פורמטים של FP8 ללמידה עמוקה.
  • הסוגים f8E4M3FNUZ ו-f8E5M2FNUZ תואמים לקידודים E4M3 ו-E5M2 של פורמטי FP8 שמתוארים במאמר 8-bit Numerical Formats for Deep Neural Networks.
  • הסוג f8E4M3B11FNUZ שמתאים לקידוד E4M3 של פורמטי FP8 שמתוארים במאמר Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks.
  • הסוג bf16 שמתאים לפורמט bfloat16 שמתואר במאמר BFloat16: The secret to high performance on Cloud TPUs.
  • הסוגים f16, f32 ו-f64 תואמים לפורמטים binary16 (דיוק חצי), binary32 (דיוק יחיד) ו-binary64 (דיוק כפול) שמתוארים בתקן IEEE 754.
  • הסוג tf32 תואם לפורמט TensorFloat32 ויש לו תמיכה מוגבלת ב-StableHLO.
  • ‫f4E2M1FN, f6E2M3FN, f6E3M2FN ו-f8E8M0FNU סוגי MX (מיקרו-סקיילינג) שמתוארים במפרט הפורמטים של OCP Microscaling.
• סוגים מורכבים מייצגים ערכים מורכבים שיש להם חלק ממשי וחלק מדומה מאותו סוג אלמנט. הסוגים המורכבים הנתמכים הם complex<f32> (שני החלקים הם מסוג f32) ו-complex<f64> (שני החלקים הם מסוג f64).

FunctionType ::= '(' InputTypes ')' '->' '(' OutputTypes ')'
InputTypes ::= [ValueType {',' ValueType}]
OutputTypes ::= [ValueType {',' ValueType}]

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

StringType ::= 'string'

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

תפעול

פעולות StableHLO (שנקראות גם ops) מייצגות קבוצה סגורה של פעולות ברמה גבוהה במודלים של למידת מכונה. כמו שצוין למעלה, התחביר של StableHLO מבוסס במידה רבה על MLIR, שהוא לא בהכרח החלופה הכי נוחה, אבל אפשר לטעון שהוא הכי מתאים למטרה של StableHLO – ליצור יכולת פעולה הדדית טובה יותר בין מסגרות ML וקומפיילרים של ML.

Op            ::= [OpOutputs] OpName OpInputs ':' OpSignature
OpName        ::= '"' 'stablehlo' '.' OpMnemonic '"'
OpMnemonic    ::= 'abs' | 'add' | ...

לפעולות StableHLO (שנקראות גם ops) יש שם, קלט/פלט וחתימה. השם מורכב מהקידומת stablehlo. ומנימוניק שמזהה באופן ייחודי אחת מהפעולות הנתמכות. בהמשך מופיעה רשימה מקיפה של כל הפעולות הנתמכות.

OpInputs        ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues   ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue    ::= ValueId
OpInputFuncs    ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs    ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs       ::= [OpOutput {',' OpOutput} '=']
OpOutput        ::= ValueId

פעולות צורכות קלטים ומפיקות פלטים. הקלט מסווג לערכי קלט (מחושבים במהלך ההפעלה), פונקציות קלט (מסופקות באופן סטטי, כי בפונקציות StableHLO לא מוגדרים ערכים מסוג first-class) ולמאפייני קלט (מסופקים גם באופן סטטי). סוגי הקלט והפלט שפעולה צורכת ומייצרת תלויים בסימן שלה. לדוגמה, הפעולה add צורכת 2 ערכי קלט ומפיקה ערך פלט אחד. לעומת זאת, הפעולה select_and_scatter צורכת 3 ערכי קלט, 2 פונקציות קלט ו-3 מאפייני קלט.

OpInputFunc ::= '{' Unused FuncInputs ':' FuncBody '}'
Unused      ::= '^' digit {digit}
              | '^' letter {letter | digit}

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

התחביר של פונקציות הקלט כולל חלק שלא נמצא בשימוש כרגע (ראו את Unusedהייצור שלמעלה), והוא נועד לתאימות ל-MLIR. ב-MLIR, יש מושג כללי יותר של 'אזורים' שיכולים לכלול כמה 'בלוקים' של פעולות שמחוברים באמצעות פעולות קפיצה. לבלוקים האלה יש מזהים שתואמים לUnused הייצור, כדי שאפשר יהיה להבדיל ביניהם. ב-StableHLO אין אופרטורים של קפיצה, ולכן החלק התואם בתחביר של MLIR לא נמצא בשימוש (אבל הוא עדיין קיים).

OpInputAttr      ::= OpInputAttrName '=' OpInputAttrValue
OpInputAttrName  ::= letter {letter | digit}
OpInputAttrValue ::= Constant

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

בשלב הזה, תוכניות StableHLO בשימוש נרחב מכילות לפעמים מאפיינים שלא מתוארים במסמך הזה. בעתיד, אנחנו מתכננים להוסיף את המאפיינים האלה ל-opset של StableHLO או לאסור את השימוש בהם בתוכניות של StableHLO. בינתיים, הנה רשימת המאפיינים האלה:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (מס' 619).
• output_operand_aliases (#740).
• מטא-נתונים של מיקום (#594).

OpSignature ::= '(' [ValueType {',' ValueType}] ')' '->' '(' [ValueType {',' ValueType}] ')'

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

בהמשך מופיעה דוגמה לאופרטור שהנימוניק שלו הוא select_and_scatter. היא צורכת 3 ערכי קלט (%operand, ‏ %source ו-%init_value), 2 פונקציות קלט ו-3 מאפייני קלט (window_dimensions, ‏ window_strides ו-padding). שימו לב שהחתימה של הפעולה כוללת רק את הסוגים של ערכי הקלט שלה (אבל לא את הסוגים של פונקציות הקלט והמאפיינים שמוגדרים בשורה).

%result = "stablehlo.select_and_scatter"(%operand, %source, %init_value) ({
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GE>
    } : (tensor<i32>, tensor<i32>) -> tensor<i1>
    "stablehlo.return"(%0) : (tensor<i1>) -> ()
}, {
  ^bb0(%arg0: tensor<i32>, %arg1: tensor<i32>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "stablehlo.return"(%0) : (tensor<i32>) -> ()
}) {
  window_dimensions = dense<[3, 1]> : tensor<2xi64>,
  window_strides = dense<[2, 1]> : tensor<2xi64>,
  padding = dense<[[0, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<4x2xi32>, tensor<2x2xi32>, tensor<i32>) -> tensor<4x2xi32>

ערכים קבועים

Constant ::= BooleanConstant
           | IntegerConstant
           | FloatConstant
           | ComplexConstant
           | TensorConstant
           | QuantizedTensorConstant
           | StringConstant
           | EnumConstant

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

BooleanConstant ::= BooleanLiteral
BooleanLiteral  ::= 'true' | 'false'

קבועים בוליאניים מייצגים ערכים בוליאניים true ו-false. לקבועים בוליאניים יש סוג i1.

IntegerConstant   ::= IntegerLiteral ':' IntegerType
IntegerLiteral    ::= ['-' | '+'] DecimalDigits
                    | ['-' | '+'] '0x' HexadecimalDigits
DecimalDigits     ::= decimalDigit {decimalDigit}
HexadecimalDigits ::= hexadecimalDigit {hexadecimalDigit}
decimalDigit      ::= '0' | ... | '9'
hexadecimalDigit  ::= decimalDigit | 'a' | ... | 'f' | 'A' | ... | 'F'

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

• (C1) is_wellformed(integer_literal, integer_type).

FloatConstant  ::= FloatLiteral ':' FloatType
FloatLiteral   ::= SignPart IntegerPart FractionalPart ScientificPart
                 | '0x' [HexadecimalDigits]
SignPart       ::= ['-' | '+']
IntegerPart    ::= DecimalDigits
FractionalPart ::= ['.' [DecimalDigits]]
ScientificPart ::= [('e' | 'E') ['-' | '+'] DecimalDigits]

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

• ‫(C1) אם משתמשים בסימון לא הקסדצימלי, is_wellformed(float_literal, float_type).
• ‫(C2) אם משתמשים בסימון הקסדצימלי, size(hexadecimal_digits) = num_bits(float_type) / 4.

ComplexConstant ::= ComplexLiteral ':' ComplexType
ComplexLiteral  ::= '(' RealPart ',' ImaginaryPart ')'
RealPart        ::= FloatLiteral
ImaginaryPart   ::= FloatLiteral

קבועים מרוכבים מייצגים ערכים מרוכבים באמצעות רשימות של חלק ממשי (מופיע ראשון) וחלק דמיוני (מופיע שני). לדוגמה, (1.0, 0.0) : complex<f32> מייצג את 1.0 + 0.0i, ו-(0.0, 1.0) : complex<f32> מייצג את 0.0 + 1.0i. הסדר שבו החלקים האלה נשמרים בזיכרון מוגדר בהטמעה. יש הגבלות על קבועים מורכבים:

• (C1) is_wellformed(real_part, complex_element_type(complex_type)).
• ‫(C2) is_wellformed(imaginary_part, complex_element_type(complex_type)).

TensorConstant ::= TensorLiteral ':' TensorType
TensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'
DenseLiteral   ::= DenseDimension | DenseElements
DenseDimension ::= '[' [DenseLiteral {',' DenseLiteral}] ']'
DenseElements  ::= [ElementLiteral {',' ElementLiteral}]
ElementLiteral ::= BooleanLiteral | IntegerLiteral | FloatLiteral | ComplexLiteral

קבועים של טנסור מייצגים ערכי טנסור באמצעות רשימות מקוננות שמוגדרות באמצעות סימון NumPy. לדוגמה, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> מייצג ערך טנזור עם המיפוי הבא מאינדקסים לרכיבים: {0, 0} => 1, ‏ {0, 1} => 2, ‏ {0, 2} => 3, ‏ {1, 0} => 4, ‏ {1, 1} => 5,‏ {1, 2} => 6. הסדר שבו הרכיבים האלה נשמרים בזיכרון מוגדר בהטמעה. ההגבלות הבאות חלות על קבועים של Tensor:

• ‫(C1) has_syntax(tensor_literal, element_type(tensor_type)), כאשר:
  • has_syntax(element_literal: Syntax, element_type: Type) = is_wellformed(element_literal, type).
  • has_syntax(tensor_literal: List, element_type: Type) = has_syntax(tensor_literal..., element_type).
• ‫(C2) has_shape(tensor_literal, shape(tensor_type)), כאשר:
  • has_shape(element_literal: Syntax, []) = true.
  • has_shape(tensor_literal: List, shape: List) = size(tensor_literal) = shape[0] and has_shape(tensor_literal..., shape[1:]).
  • אחרת, false.

QuantizedTensorConstant ::= QuantizedTensorLiteral ':' QuantizedTensorType
QuantizedTensorLiteral  ::= 'dense' '<' (DenseLiteral | ElementLiteral) '>'

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

• (C1) has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type)).
• ‫(C2) has_shape(quantized_tensor_literal, shape(quantized_tensor_type)).

StringConstant  ::= StringLiteral
StringLiteral   ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence  ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))

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

תפעול

מוחלט

סמנטיקה

מבצע פעולת abs לפי רכיבים בטנזור operand ומפיק טנזור result. בהתאם לסוג הרכיב, מבצעים את הפעולות הבאות:

• למספרים שלמים עם סימן: מודולוס של מספר שלם.
• למספרים ממשיים: abs מ-IEEE-754.
• למספרים מרוכבים: מודולוס מרוכב.
• לסוגים כמותיים: dequantize_op_quantize(abs, operand, type(result)).

קלט

פלט

מגבלות

• (C1) shape(result) = shape(operand).
• ‫(C2) baseline_element_type(result) מוגדר כך:
  • ‫complex_element_type(element_type(operand)) if is_complex(operand).
  • baseline_element_type(operand) אחרת.

דוגמאות

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand)< : (t>ens>or3xi32<) - t>ensor3xi32
// %result: [2, 0, 2]

 דוגמאות נוספות

הוספה

סמנטיקה

מבצעת חיבור של שני טנסורים lhs ו-rhs לפי רכיבים ומפיקה טנסור result. בהתאם לסוג הרכיב, מבצעים את הפעולות הבאות:

• לערכים בוליאניים: OR לוגי.
• למספרים שלמים: חיבור של מספרים שלמים.
• למספרים ממשיים: addition מ-IEEE-754.
• למספרים מרוכבים: חיבור מרוכב.
• לסוגים כמותיים: dequantize_op_quantize(add, lhs, rhs, type(result)).

קלט

פלט

מגבלות

• אם הפעולה משתמשת בטנסורים לא מכומתים:
  • (C1) type(lhs) = type(rhs) = type(result).
• אם הפעולה משתמשת בטנסורים שעברו קוונטיזציה:
  • ‫(C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
  • ‫(C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • ‫(C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • ‫(C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • ‫(C6) אם is_per_axis_quantized(lhs), אז quantization_dimension(lhs) = quantization_dimension(result).
  • ‫(C7) אם is_per_axis_quantized(rhs), אז quantization_dimension(rhs) = quantization_dimension(result).

דוגמאות

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.add"(%lhs, %rhs)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[6, 8], [10, 12]]

 דוגמאות נוספות

after_all

סמנטיקה

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

קלט

פלט

דוגמאות

// %input0: !stablehlo.token
// %input1: !stablehlo.token
%result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehl>o.token) - !stablehlo.token

 דוגמאות נוספות

all_gather

סמנטיקה

בכל קבוצת תהליכים ברשת התהליכים של StableHLO, הפונקציה מחברת את הערכים של טנסורים operands מכל תהליך לאורך all_gather_dim ומפיקה טנסורים results.

הפעולה מפצלת את רשת התהליכים של StableHLO ל-process_groups, שמוגדרת כך:

• cross_replica(replica_groups) if channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups) if channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups) if channel_id > 0 and use_global_device_ids = true.

לאחר מכן, בכל process_group:

• ‫operands...@receiver = [operand@sender for sender in process_group] לכל receiver ב-process_group.
• ‫results...@process = concatenate(operands...@process, all_gather_dim) לכל process ב-process_group.

קלט

פלט

מגבלות

• (C1) 0 <= all_gather_dim < rank(operands...).
• ‫(C2) is_unique(replica_groups).
• ‫(C3) size(replica_groups) מוגדר כך:
  • ‫num_replicas אם משתמשים ב-cross_replica.
  • ‫num_replicas אם משתמשים ב-cross_replica_and_partition.
  • ‫num_processes אם משתמשים ב-flattened_ids.
• ‫(C4) 0 <= replica_groups < size(replica_groups).
• ‫(C5) אם use_global_device_ids = true, אז channel_id > 0.
• ‫(C6) type(results...) = type(operands...) except:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

דוגמאות

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64,
  // channel_id = 0
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
  // use_global_device_ids = false
}< : (ten>sor2x2xi<64, ten>sor>2x2xi64)< - (ten>sor2x4xi<64, ten>sor2x4xi64)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

 דוגמאות נוספות

all_reduce

סמנטיקה

בתוך כל קבוצת תהליכים ברשת התהליכים של StableHLO, הפונקציה מחילה צמצום computation על ערכי הטנסורים operands מכל תהליך ומפיקה טנסורים results.

הפעולה מפצלת את רשת התהליכים של StableHLO ל-process_groups, שמוגדרת כך:

• cross_replica(replica_groups) if channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups) if channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups) if channel_id > 0 and use_global_device_ids = true.

לאחר מכן, בכל process_group:

• ‫results...@process[result_index] = exec(schedule) עבור עץ בינארי מסוים ‫schedule כאשר:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• ‫schedule הוא עץ בינארי שמוגדר בהטמעה, והמעבר שלו בסדר עולה הוא to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

קלט

פלט

מגבלות

• (C1) is_unique(replica_groups).
• ‫(C2) size(replica_groups) מוגדר כך:
  • ‫num_replicas אם משתמשים ב-cross_replica.
  • ‫num_replicas אם משתמשים ב-cross_replica_and_partition.
  • ‫num_processes אם משתמשים ב-flattened_ids.
• ‫(C3) 0 <= replica_groups < size(replica_groups).
• ‫(C4) אם use_global_device_ids = true, אז channel_id > 0.
• ‫(C5) computation הוא מסוג (tensor<E>, tensor<E>) -> (tensor<E>) כאשר is_promotable(element_type(operand), E).
• ‫(C6) shape(results...) = shape(operands...).
• ‫(C7) element_type(results...) = E.

דוגמאות

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%ar<g0:> tensori64, %ar<g1:> tensori64):
    %0 = "stablehlo.add"(%ar<g0,> %arg1) <: (>ten>sori64,< te>nsori64) - tensori64
    "stable<hlo>.re>turn"(%0) : (tensori64) - ()<
}) {
  >replica_g<roups => dense[[0, 1]] : tensor1x2xi64,
  // channel_id = 0
  channel_hand<le = #stablehlo.chan>nel_handlehandle = 0, type = 0
  // use_global_<devic>e_ids = <false>
} >: (tenso<r4xi6>4, tenso<r4xi6>4) - (tensor4xi64, tensor4xi64)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

 דוגמאות נוספות

all_to_all

סמנטיקה

בכל קבוצת תהליכים ברשת התהליכים של StableHLO, הפונקציה מפצלת את הערכים של טנסורים operands לאורך split_dimension לחלקים, מפזרת את החלקים המפוצלים בין התהליכים, משרשרת את החלקים המפוזרים לאורך concat_dimension ומפיקה טנסורים results. הפעולה מפצלת את רשת התהליכים של StableHLO ל-process_groups, שמוגדרת כך:

• ‫cross_replica(replica_groups) if channel_id <= 0.
• ‫cross_partition(replica_groups) if channel_id > 0.

לאחר מכן, בכל process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) לכל sender ב-process_group.
• ‫scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] כשreceiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

קלט

פלט

מגבלות

• (C1) 0 <= split_dimension < rank(operands...).
• ‫(C2) dim(operands..., split_dimension) % split_count = 0.
• ‫(C3) 0 <= concat_dimension < rank(operands...).
• ‫(C4) 0 < split_count.
• ‫(C5) is_unique(replica_groups).
• ‫(C6) size(replica_groups) מוגדר כך:
  • ‫num_replicas אם משתמשים ב-cross_replica.
  • ‫num_partitions אם משתמשים ב-cross_partition.
• ‫(C7) 0 <= replica_groups < size(replica_groups).
• ‫(C8) dim(replica_groups, 1) = split_count.
• ‫(C9) type(results...) = type(operands...) למעט, אם split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

דוגמאות

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_grou<ps = den>se[[0, 1]<] : ten>sor1x2xi64
  // channel_id = 0
}< : (ten>sor2x4xi<64, ten>sor>2x4xi64)< - (ten>sor4x2xi<64, ten>sor4x2xi64)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

 דוגמאות נוספות

וגם

סמנטיקה

מבצעת פעולת AND לפי רכיבים של שני טנסורים, lhs ו-rhs, ומפיקה טנסור result. בהתאם לסוג הרכיב, מבצעים את הפעולות הבאות:

• לערכים בוליאניים: AND לוגי.
• למספרים שלמים: ערך AND ברמת הביטים.

קלט

פלט

מגבלות

• (C1) type(lhs) = type(rhs) = type(result).

דוגמאות

// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.and"(%lhs, %rhs)< : (ten>sor2x2xi<32, ten>sor>2x2xi32<) - ten>sor2x2xi32
// %result: [[1, 2], [3, 0]]

 דוגמאות נוספות

atan2

סמנטיקה

מבצע פעולת atan2 לפי רכיבים בטנזורים lhs ו-rhs ומפיק טנזור result. בהתאם לסוג הרכיב, מבצעים את הפעולות הבאות:

• למספרים ממשיים: atan2 מ-IEEE-754.
• למספרים מרוכבים: complex atan2.
• לסוגים כמותיים: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

קלט

פלט

מגבלות

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result).

דוגמאות

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs)< : (t>ensor3xf<64, t>ens>or3xf64<) - t>ensor3xf64
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

 דוגמאות נוספות

batch_norm_grad

סמנטיקה

מחשב את הגרדיאנטים של כמה תשומות של batch_norm_training באמצעות backpropagation מ-grad_output, ומפיק טנסורים של grad_operand, grad_scale ו-grad_offset. באופן רשמי יותר, אפשר לבטא את הפעולה הזו כפירוק לפעולות קיימות של StableHLO באמצעות תחביר Python באופן הבא:

def compute_sum(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  return sum

def compute_mean(operand, feature_index):
  sum = compute_sum(operand, feature_index)
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
  # Broadcast inputs to type(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance`
  # Intermediate values will be useful for computing gradients
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)

  # Use the implementation from batchnorm_expander.cc in XLA
  # Temporary variables have exactly the same names as in the C++ code
  elements_per_feature = broadcast_in_dim(
      constant(divide(size(operand), dim(operand, feature_index)),
               element_type(grad_output)),
      [], type(operand))
  i1 = multiply(grad_output, elements_per_feature)
  i2 = broadcast_in_dim(
      compute_sum(grad_output, feature_index), [feature_index], type(operand))
  i3 = broadcast_in_dim(
      compute_sum(multiply(grad_output, centered_operand), feature_index),
      [feature_index], type(operand))
  i4 = multiply(i3, centered_operand)
  i5 = divide(i4, add(variance_bcast, epsilon_bcast))
  i6 = subtract(subtract(i1, i2), i5)

  grad_operand =
      multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
  grad_scale =
      compute_sum(multiply(grad_output, normalized_operand), feature_index)
  grad_offset = compute_sum(grad_output, feature_index)

  return grad_operand, grad_scale, grad_offset

לסוגים כמותיים, הפונקציה מבצעת dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index)).

קלט

פלט

מגבלות

• (C1) 0 <= feature_index < rank(operand).
• ‫(C2) operand, ‏ scale, ‏ mean, ‏ variance, ‏ grad_output, ‏ grad_operand,‏ grad_scale ו-grad_offset כוללים את אותו baseline_element_type.
• ‫(C3) ל-operand, ל-grad_output ול-grad_operand יש את אותו הצורה.
• ‫(C4) ל-scale, ‏ mean, ‏ variance, ‏ grad_scale ו-grad_offset יש את אותה צורה.
• ‫(C5) size(scale) = dim(operand, feature_index).

דוגמאות

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf64,
 <    tenso>r2x>2x2xf64)< - (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

סמנטיקה

מבצעת נורמליזציה של טנסור operand בכל המימדים, למעט מימד feature_index, ומפיקה טנסור result. באופן רשמי יותר, אפשר לבטא את הפעולה הזו כפירוק לפעולות קיימות של StableHLO באמצעות תחביר Python באופן הבא:

def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
  # Broadcast inputs to shape(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance` instead of
  # computing them like `batch_norm_training` does.
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)
  return add(multiply(scale_bcast, normalized_operand), offset_bcast)

לסוגים כמותיים, הפונקציה מבצעת dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result)).

קלט

פלט

מגבלות

• (C1) 0 <= feature_index < rank(operand).
• ‫(C2) operand, scale, offset, mean, variance ו-result כולן בעלות אותו baseline_element_type.
• ‫(C3) size(scale) = dim(operand, feature_index).
• ‫(C4) size(offset) = dim(operand, feature_index).
• ‫(C5) size(mean) = dim(operand, feature_index).
• ‫(C6) size(variance) = dim(operand, feature_index).
• ‫(C7) baseline_type(operand) = baseline_type(result).

דוגמאות

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ensor2xf<64, t>ens>or2xf64<) - tenso>r2x2x2xf64
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

סמנטיקה

מחשבת את הממוצע והשונות בכל המאפיינים, למעט המאפיין feature_index ומנרמלת את טנסור operand כדי ליצור את הטנסורים output, batch_mean ו-batch_var. באופן רשמי יותר, אפשר לבטא את הפעולה הזו כפירוק לפעולות קיימות של StableHLO באמצעות תחביר Python באופן הבא:

def compute_mean(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def compute_variance(operand, feature_index):
  mean = compute_mean(operand, feature_index)
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  centered_operand = subtract(operand, mean_bcast)
  return compute_mean(mul(centered_operand, centered_operand), feature_index)

def batch_norm_training(operand, scale, offset, epsilon, feature_index):
  mean = compute_mean(operand, feature_index)
  variance = compute_variance(operand, feature_index)
  return batch_norm_inference(operand, scale, offset, mean, variance, epsilon,
                              feature_index),
         mean, variance

לסוגים כמותיים, הפונקציה מבצעת dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var)).

קלט

פלט

מגבלות

• (C1) 0 <= feature_index < rank(operand).
• ‫(C2) operand, ‏ scale, ‏ offset, ‏ batch_mean, ‏ batch_var ו-output כוללים את אותו baseline_element_type.
• ‫(C3) size(scale) = dim(operand, feature_index).
• ‫(C4) size(offset) = dim(operand, feature_index).
• ‫(C5) size(batch_mean) = dim(operand, feature_index).
• ‫(C6) size(batch_var) = dim(operand, feature_index).
• ‫(C7) baseline_type(output) = baseline_type(operand).

דוגמאות

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
}< : (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ens>or2xf64) -
 <   (tenso>r2x2x2xf<64, t>ensor2xf<64, t>ensor2xf64)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

סמנטיקה

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

באופן רשמי יותר, בהינתן E = element_type(operand), E' = element_type(result), ו-R = rank(operand):

• אם num_bits(E') < num_bits(E), bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]).
• אם num_bits(E') > num_bits(E), bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :]).
• אם num_bits(E') = num_bits(E), bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1]).

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

קלט

פלט

מגבלות

• ‫(C1) בהינתן E = is_quantized(operand) ? storage_type(operand) : element_type(operand), ‏E' = is_quantized(result) ? storage_type(result) : element_type(result) ו-R = rank(operand):
  • אם num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • אם num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) לכל 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • אם num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) לכל 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• ‫(C2) אם is_complex(operand) or is_complex(result), אז is_complex(operand) and is_complex(result).

דוגמאות

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand)< : >(te>nsorf64<) - t>ensor4xf16
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

 דוגמאות נוספות

broadcast_in_dim

סמנטיקה

מרחיב את המאפיינים או את הדירוג של טנסור קלט על ידי שכפול הנתונים בטנסור operand ומפיק טנסור result. באופן רשמי יותר, result[result_index] = operand[operand_index] כאשר לכל d ב- axes(operand):

• ‫operand_index[d] = 0 if dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]] אחרת.

קלט

פלט

מגבלות

• ‫(C1) element_type(result) מחושב לפי הנוסחה הבאה:
  • ‫element_type(operand), אם !is_per_axis_quantized(operand).
  • ‫element_type(operand), למעט quantization_dimension(operand),‏ scales(operand) ו-zero_points(operand) שעשויים להיות שונים מ-quantization_dimension(result),‏ scales(result) ו-zero_points(result) בהתאמה, אחרת.
• ‫(C2) size(broadcast_dimensions) = rank(operand).
• ‫(C3) 0 <= broadcast_dimensions < rank(result).
• ‫(C4) is_unique(broadcast_dimensions).
• ‫(C5) לכל d בaxes(operand):
  • dim(operand, d) = 1 או
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• ‫(C6) אם is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • אם dim(operand, quantization_dimension(operand)) = 1, אז scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

דוגמאות

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensio<ns = arra>yi64: 2, 1
}< : (ten>sor>1x3xi32<) - tenso>r2x3x2xi32
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 דוגמאות נוספות

כיסוי

סמנטיקה

הפונקציה מפיקה את הפלט מהרצה של פונקציה אחת בדיוק מתוך branches בהתאם לערך של index. באופן רשמי יותר, result = selected_branch() כאשר:

• ‫selected_branch = branches[index] if 0 <= index < size(branches).
• selected_branch = branches[-1] אחרת.

קלט

פלט

מגבלות

• (C1) 0 < size(branches).
• ‫(C2) input_types(branches...) = [].
• ‫(C3) same(output_types(branches...)).
• ‫(C4) type(results...) = output_types(branches[0]).

דוגמאות

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %resul<t_bra>nch0) : <(tens>or2>xi64, tensor2xi64) - ()
}, {
  "stablehlo.return"(%result_branc<h1, %>result_b<ranch>1) >: (tensor2xi64, <ten>sor>2xi64) -< ()
}>) : (ten<sori3>2) - (tensor2xi64, tensor2xi64)
// %result0: [1, 1]
// %result1: [1, 1]

 דוגמאות נוספות

cbrt

סמנטיקה

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

• למספרים ממשיים: rootn(x, 3) מ-IEEE-754.
• למספרים מרוכבים: שורש מעוקב מרוכב.
• לסוגים שעברו קוונטיזציה: dequantize_op_quantize(cbrt, operand, type(result))

קלט

פלט

מגבלות

• (C1) baseline_type(operand) = baseline_type(result).

דוגמאות

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand)< : (t>ens>or4xf64<) - t>ensor4xf64
// %result: [0.0, 1.0, 2.0, 3.0]

 דוגמאות נוספות

ceil

סמנטיקה

מבצעת פעולת ceil על כל רכיב בטנזור operand ומפיקה טנזור result. מיישמת את הפעולה roundToIntegralTowardPositive מהמפרט IEEE-754. לסוגים כמותיים, הפונקציה מבצעת dequantize_op_quantize(ceil, operand, type(result)).

קלט

פלט

מגבלות

• (C1) baseline_type(operand) = baseline_type(result).

דוגמאות

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.ceil"(%operand)< : (t>ens>or5xf32<) - t>ensor5xf32
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

 דוגמאות נוספות

cholesky

סמנטיקה

מחשבת את פירוק Cholesky של קבוצת מטריצות.

באופן רשמי יותר, לכל i ב-index_space(result),‏ result[i0, ..., iR-3, :, :] הוא פירוק שולסקי של a[i0, ..., iR-3, :, :], בצורה של מטריצה משולשת תחתונה (אם lower הוא true) או מטריצה משולשת עליונה (אם lower הוא false). ערכי הפלט במשולש הנגדי, כלומר במשולש העליון או במשולש התחתון, בהתאמה, מוגדרים על ידי ההטמעה.

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

לסוגים כמותיים, הפונקציה מבצעת dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

קלט

פלט

מגבלות

• (C1) baseline_type(a) = baseline_type(result).
• ‫(C2) 2 <= rank(a).
• ‫(C3) dim(a, -2) = dim(a, -1).

דוגמאות

// %a: [
//      [1.0, 2.0, 3.0],
//      [2.0, 20.0, 26.0],
//      [3.0, 26.0, 70.0]
//     ]
%result = "stablehlo.cholesky"(%a) {
  lower = true
}< : (ten>sor>3x3xf32<) - ten>sor3x3xf64
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

תיחום

סמנטיקה

הפונקציה מחזירה טנזור result שבו כל רכיב בטנזור operand מוגבל לערך מינימלי ומקסימלי. באופן רשמי יותר, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), כאשר min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. לסוגים כמותיים, הפונקציה מבצעת dequantize_op_quantize(clamp, min, operand, max, type(result)).

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

קלט

פלט

מגבלות

• (C1) rank(min) = 0 or shape(min) = shape(operand).
• ‫(C2) rank(max) = 0 or shape(max) = shape(operand).
• ‫(C3) baseline_element_type(min) = baseline_element_type(operand) = baseline_element_type(max).
• ‫(C4) baseline_type(operand) = baseline_type(result).

דוגמאות

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max)< : (t>ensor3xi<32, t>ensor3xi<32, t>ens>or3xi32<) - t>ensor3xi32
// %result: [5, 13, 20]

 דוגמאות נוספות

collective_broadcast

סמנטיקה

בכל קבוצת תהליכים ברשת התהליכים של StableHLO, הערך של טנזור operand נשלח מתהליך המקור לתהליכי היעד ונוצר טנזור result.

הפעולה מפצלת את רשת התהליכים של StableHLO ל-process_groups, שמוגדרת כך:

• ‫cross_replica(replica_groups) if channel_id <= 0.
• ‫cross_partition(replica_groups) if channel_id > 0.

לאחר מכן, result@process מחושב לפי הנוסחה הבאה:

• ‫operand@process_groups[i, 0] אם קיים i כך שהתהליך נמצא ב-process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) אחרת.

קלט

פלט

מגבלות

• (C1) is_unique(replica_groups).
• ‫(C2) 0 <= replica_groups < N כאשר N מוגדר כך:
  • ‫num_replicas אם משתמשים ב-cross_replica.
  • ‫num_partitions אם משתמשים ב-cross_partition.
• ‫(C3) type(result) = type(operand).

דוגמאות

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_grou<ps = den>se[[2, 1]<] : ten>sor1x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
} : (ten>sor>1x2xi64<) - ten>sor1x2xi64
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

סמנטיקה

בכל קבוצת תהליכים ברשת התהליכים של StableHLO, הערך של טנזור operand נשלח מתהליך המקור לתהליך היעד ונוצר טנזור result.

הפעולה מפצלת את רשת התהליכים של StableHLO ל-process_groups, שמוגדרת כך:

• ‫cross_replica(source_target_pairs) if channel_id <= 0.
• ‫cross_partition(source_target_pairs) if channel_id > 0.

לאחר מכן, result@process מחושב לפי הנוסחה הבאה:

• ‫operand@process_groups[i, 0], אם קיים i כך ש- process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) אחרת.

קלט

פלט

מגבלות

• (C1) dim(source_target_pairs, 1) = 2.
• ‫(C2) is_unique(source_target_pairs[:, 0]).
• ‫(C3) is_unique(source_target_pairs[:, 1]).
• ‫(C4) 0 <= source_target_pairs < N, כאשר N מוגדר כך:
  • ‫num_replicas אם משתמשים ב-cross_replica.
  • ‫num_partitions אם משתמשים ב-cross_partition.
• ‫(C5) type(result) = type(operand).

דוגמאות

// num_replicas: 3
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
// %operand@(2, 0): [[9, 10], [11, 12]]
%result = "stablehlo.collective_permute"(%operand) {
  source_target_pai<rs = dense[[0, 1>], [1, 2]<] : ten>sor2x2xi64,
  channel_handle = #stablehlo.chan<nel_handlehandle = 0>, type = 0
}< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
//
// %result@(0, 0): [[0, 0], [0, 0]]
// %result@(1, 0): [[1, 2], [3, 4]]
// %result@(2, 0): [[5, 6], [7, 8]]

 דוגמאות נוספות

השוואה

סמנטיקה

מבצע השוואה בין רכיבים של טנסורים lhs ו-rhs בהתאם ל-comparison_direction ול-compare_type, ומפיק טנסור result.

המשמעות של הערכים comparison_direction ו-compare_type היא כדלקמן:

לסוגי רכיבים בוליאניים ומספריים שלמים:

• EQ: lhs = rhs.
• NE: lhs != rhs.
• GE: lhs >= rhs.
• GT: lhs > rhs.
• LE: lhs <= rhs.
• LT: lhs < rhs.

עבור סוגי רכיבים של נקודה צפה עם compare_type = FLOAT, האופרטור מטמיע את הפעולות הבאות של IEEE-754:

• EQ: compareQuietEqual.
• NE: compareQuietNotEqual.
• GE: compareQuietGreaterEqual.
• GT: compareQuietGreater.
• LE: compareQuietLessEqual.
• LT: compareQuietLess.

עבור סוגי רכיבים של נקודה צפה עם compare_type = TOTALORDER, האופרטור משתמש בשילוב של פעולות totalOrder ו-compareQuietEqual מ-IEEE-754.

לסוגי רכיבים מורכבים, מתבצעת השוואה לקסיקוגרפית של זוגות (real, imag) באמצעות comparison_direction ו-compare_type שסופקו. הגדרת סדר למספרים מרוכבים כוללת סמנטיקה מפתיעה, ולכן בעתיד אנחנו מתכננים להסיר את התמיכה במספרים מרוכבים כש-comparison_direction הוא GE, ‏ GT, ‏ LE או LT (מספר 560).

לסוגים כמותיים. הפונקציה מבצעת dequantize_compare(lhs, rhs, comparison_direction).

קלט

פלט

מגבלות

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• ‫(C2) shape(lhs) = shape(rhs) = shape(result).
• ‫(C3) compare_type מוגדר כך:
  • ‫SIGNED if is_signed_integer(element_type(lhs)).
  • ‫UNSIGNED if is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • ‫FLOAT או TOTALORDER אם is_float(element_type(lhs)).
  • ‫FLOAT if is_complex(element_type(lhs)).

דוגמאות

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = <#stablehlocomparison_di>rection LT,
  compare_type = <#stablehlocomparison_>type FLOAT
}< : (t>ensor2xf<32, t>ens>or2xf32<) - >tensor2xi1
// %result: [true, false]

 דוגמאות נוספות

מורכב

סמנטיקה

מבצעת המרה של כל רכיב בנפרד לערך מרוכב מתוך זוג של ערכים ממשיים ומדומים, lhs ו-rhs, ומפיקה טנזור result.

קלט

פלט

מגבלות

• (C1) type(lhs) = type(rhs).
• ‫(C2) shape(result) = shape(lhs).
• ‫(C3) element_type(result) הוא מסוג complex<E> כאשר E = element_type(lhs).

דוגמאות

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs)< : (t>ensor2xf<64, t>ens>or2xf64<) - tenso<r2x>>complexf64
// %result: [(1.0, 2.0), (3.0, 4.0)]

 דוגמאות נוספות

מורכב

סמנטיקה

פעולה שמורכבת מפעולות אחרות של StableHLO, מקבלת inputs ו-composite_attributes ומפיקה results. הסמנטיקה של הפעולה מיושמת באמצעות המאפיין decomposition. אפשר להחליף את הפעולה composite בפירוק שלה בלי לשנות את הסמנטיקה של התוכנית. במקרים שבהם הוספת הפירוק לשורה לא מספקת את אותה סמנטיקה של פעולה, עדיף להשתמש ב-custom_call.

השדה version (ברירת המחדל היא 0) משמש לציון המועד שבו הסמנטיקה של רכיב מורכב משתנה.

קלט

פלט

מגבלות

• ‫(C1) is_namespaced_op_name(name)
• ‫(C2) is_defined_in_parent_scope(decomposition)
• ‫(C3) types(inputs...) == input_types(decomposition)
• ‫(C4) types(results...) == output_types(decomposition)

דוגמאות

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
 < ve>rsion = <1 :> i3>2
} : (<ten>sorf32, tensorf32) - tensorf32

 דוגמאות נוספות

concatenate

סמנטיקה

משרשר את inputs לאורך המימד dimension באותו סדר של הארגומנטים שצוינו ומפיק טנסור result. בצורה רשמית יותר, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], כאשר:

1. id = d0 + ... + dk-1 + kd.
2. ‫d שווה ל-dimension, ו-d0, ... הם גדלי המאפיינים d של inputs.

קלט

פלט

מגבלות

• (C1) same(element_type(inputs...)).
• ‫(C2) same(shape(inputs...)) למעט dim(inputs..., dimension).
• ‫(C3) 0 < size(inputs).
• ‫(C4) 0 <= dimension < rank(inputs[0]).
• ‫(C5) element_type(result) = element_type(inputs[0]).
• ‫(C6) shape(result) = shape(inputs[0]) except for:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

דוגמאות

// %input0: [[1, 2], [3, 4], [5, 6]]
// %input1: [[7, 8]]
%result = "stablehlo.concatenate"(%input0, %input1) {
  dimension = 0 : i64
}< : (ten>sor3x2xi<64, ten>sor>1x2xi64<) - ten>sor4x2xi64
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

 דוגמאות נוספות

קבוע

סמנטיקה

הפונקציה יוצרת טנסור output מקבוע value.

קלט

פלט

מגבלות

• (C1) type(value) = type(output).

דוגמאות

%output = "stablehlo.constant"() {
  val<ue = dense[[0.0, 1.0], [>2.0, 3.0]<] : ten>sor2x2xf3>2
} : (<) - ten>sor2x2xf32
// %output: [[0.0, 1.0], [2.0, 3.0]]

 דוגמאות נוספות

להשלים המרה

סמנטיקה

מבצעת המרה מטיפוס אלמנט אחד לטיפוס אלמנט אחר בטנזור operand ומפיקה טנזור result.

בהמרות boolean-to-any-supported-type, הערך false מומר לאפס והערך true מומר לאחד. בהמרות של any-supported-type-to-boolean, ערך אפס מומר ל-false, וערכים שאינם אפס מומרים ל-true. בהמשך מוסבר איך זה עובד עם סוגים מורכבים.

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

בהמרות שכוללות floating-point-to-integer, החלק השברי מצומצם. אם אי אפשר לייצג את הערך שקוצץ בסוג היעד, ההתנהגות תוגדר בהמשך (#180).

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

בהמרות מסוג complex-to-any-other-type ו-any-other-type-to-complex, המערכת מתעלמת מהערך הדמיוני של המקור או מאפסת את הערך הדמיוני של היעד, בהתאמה. ההמרה של החלק האמיתי מתבצעת לפי המרות של נקודה צפה.

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

קלט

פלט

מגבלות

• (C1) shape(operand) = shape(result).

דוגמאות

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand)< : (t>ens>or3xi64<) - tenso<r3x>>complexf64
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

 דוגמאות נוספות

קונבולוציה

סמנטיקה

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

באופן רשמי יותר, אפשר להגדיר מחדש את הקלט במונחים של lhs כדי להגדיר חלונות של lhs:

• lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension)).
• lhs_window_strides = lhs_shape(1, window_strides, 1).
• lhs_padding = lhs_shape([0, 0], padding, [0, 0]).
• lhs_base_dilations = lhs_shape(1, lhs_dilation, 1).
• lhs_window_dilations = lhs_shape(1, rhs_dilation, 1).

השינוי הזה מתבצע באמצעות פונקציות העזר הבאות:

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]).
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]).
• ‫permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1] כשj[d] = i[permutation[d]].

אם feature_group_count = 1 ו-batch_group_count = 1, אז לכל output_spatial_index ב-index_space(dim(result, output_spatial_dimensions...)),‏ result[result_shape(:, output_spatial_index, :)] = dot_product כאשר:

• padding_value = constant(0, element_type(lhs)).
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1).
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides.
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations).
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]). נראה שהתכונה הזו לא נמצאת בשימוש, ולכן אנחנו מתכננים להסיר אותה בעתיד (מס' 1181).
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]).

אם feature_group_count > 1:

• lhses = split(lhs, feature_group_count, input_feature_dimension).
• rhses = split(rhs, feature_group_count, kernel_output_feature_dimension).
• results... = convolution(lhses..., rhses..., ..., feature_group_count=1, ...).
• result = concatenate(results, output_feature_dimension).

אם batch_group_count > 1:

• lhses = split(lhs, batch_group_count, input_batch_dimension).
• rhses = split(rhs, batch_group_count, kernel_output_feature_dimension).
• results... = convolution(lhses..., rhses..., ..., batch_group_count=1, ...).
• result = concatenate(results, output_feature_dimension).

לסוגים כמותיים, הפונקציה מבצעת dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result)).

עבור סוגים היברידיים של קוונטיזציה, הפעולה היא hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs).

קלט

פלט

מגבלות

• (C1) N = rank(lhs) = rank(rhs).
• ‫(C2) size(window_strides) = N - 2.
• ‫(C3) 0 < window_strides.
• ‫(C4) shape(padding) = [N - 2, 2].
• ‫(C5) size(lhs_dilation) = N - 2.
• ‫(C6) 0 < lhs_dilation.
• ‫(C7) size(rhs_dilation) = N - 2.
• ‫(C8) 0 < rhs_dilation.
• ‫(C9) size(window_reversal) = N - 2.
• ‫(C10) dim(lhs, input_batch_dimension) % batch_group_count = 0.
• ‫(C11) dim(lhs, input_feature_dimension) % feature_group_count = 0.
• (C12) size(input_spatial_dimensions) = N - 2.
• ‫(C13) בהינתן input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]:
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count.
• ‫(C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• ‫(C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• ‫(C17) size(kernel_spatial_dimensions) = N - 2.
• ‫(C18) בהינתן kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• ‫(C19) size(output_spatial_dimensions) = N - 2.
• ‫(C20) בהינתן output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]:
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• ‫(C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count.
• ‫(C23) feature_group_count = 1 or batch_group_count = 1.
• ‫(C24) size(precision_config) = 2.
• ‫(C25) dim(result, result_dim) מוגדר כך:
  • ‫dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension.
  • ‫dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension.
  • num_windows אחרת, כאשר:
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• אם הפעולה משתמשת בטנסורים לא מכומתים:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• אם הפעולה משתמשת בטנסורים שעברו קוונטיזציה:
  • ‫(C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • ‫(C29) אם is_per_axis_quantized(rhs), אז quantization_dimension(rhs) = kernel_output_feature_dimension.
  • ‫(C30) אם is_per_axis_quantized(result), אז quantization_dimension(result) = output_feature_dimension.
  • אם is_quantized(lhs):
  • ‫(C31) storage_type(lhs) = storage_type(rhs).
  • ‫(C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • ‫(C33) אם is_per_tensor_quantized(rhs), אז is_per_tensor_quantized(result).
  • אם !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

דוגמאות

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strid<es = arra>yi64: 4, 4,
  paddi<n>g = dense<0 : ten>sor2x2xi64,
  lhs_dilati<on = arra>yi64: 2, 2,
  rhs_dilati<on = arra>yi64: 1, 1,
  window_revers<al = arrayi1: fa>lse, false,
  // In the StableHLO dialect, dimension numbers are encoded vi<a:
  // `[input >dim<ensions]x[kernel >di>mensions]-[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" a<re spatial dimensions.
  d>imension_num>bers = #stablehlo.conv[b, 0, 1, f]x[0, 1, i, o]-[b, 0, 1, f],
  batch_group_count = 1 : i64,
  fea<ture_group_count >= 1 : i64,
 < precision_config> = [#stablehl<oprecision >DEFAULT,< #stablehlo>pre>cision <DEFAULT]
} >: (tensor1x4x4x1xi64, tensor3x3x1x1xi64) - tensor1x2x2x1xi64
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 דוגמאות נוספות

קוסינוס

סמנטיקה

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

• למספרים ממשיים: cos מ-IEEE-754.
• למספרים מרוכבים: קוסינוס מרוכב.
• לסוגים כמותיים: dequantize_op_quantize(cosine, operand, type(result)).

קלט

פלט

מגבלות

• (C1) baseline_type(operand) = baseline_type(result).

דוגמאות

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.cosine"(%operand)< : (ten>sor>2x2xf32<) - ten>sor2x2xf32
// %result: [[1.0, 0.0], [-1.0, 0.0]]

 דוגמאות נוספות

count_leading_zeros

סמנטיקה

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

קלט

פלט

מגבלות

• (C1) type(operand) = type(result).

דוגמאות

// %operand: [[0, 1], [128, -1]]
%result = "stablehlo.count_leading_zeros"(%operand)< : (ten>sor>2x2xi64<) - ten>sor2x2xi64
// %result: [[64, 63], [56, 0]]

 דוגמאות נוספות

custom_call

סמנטיקה

מכיל פעולה שמוגדרת בהטמעה call_target_name שמקבלת את inputs ואת called_computations ומפיקה את results. אפשר להשתמש ב-has_side_effect,‏ backend_config ו-api_version כדי לספק מטא-נתונים נוספים שמוגדרים על ידי ההטמעה.

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

קלט

פלט

‫(XLA GPU Support) Special custom_call targets

יש שלושה סוגים מיוחדים של call_target_name שקשורים ל-buffer: ‫CreateBuffer יוצר buffer לא מאותחל, Pin יוצר buffer מאותחל ו-Unpin מבטל את ההקצאה של buffer ומחזיר את התוכן של buffer.

%uninitialized_buffer = "stablehlo.custom_call"() {
  call_target_name = "CreateBuffer",
  api_version> = 4 : <i32,
>} : () - memref4xf64

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin&quo<t;,
 > ap>i_versi<on = >4 : i32,
} : (tensor4xf64) - memref4xf64

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_na<me = >&qu>ot;Unpi<n&quo>t;,
  api_version = 4 : i32,
} : (memref4xf64) - tensor4xf64

כינוי

יכול להיות שחלק מהפעולות של custom_call ידרשו שחלק מהפלט וחלק מהאופרנדים ישתפו את אותו הזיכרון. אפשר לבטא את זה באמצעות output_operand_aliases. ייצוג של זוג כינויים מורכב מרשימה של אינדקסים של טאפלים של פלט שמייצגים את חלק הפלט, ומ-operand_index יחד עם רשימה של אינדקסים של טאפלים של אופרנד שמייצגים את חלק האופרנד. רשימת הפלט או המדדים של טופל האופרנד ריקה אם הסוג המתאים הוא לא סוג tuple, והיא יכולה להיות ארוכה ככל שרוצים עבור סוג טופל מקונן. הייצוג הזה דומה לייצוג הכינוי XLA.

הסוג של חלק הפלט וחלק הקלט בצמד כינויים חייב להיות זהה. בפעולות custom_call שאינן קריאה אל CreateBuffer, Pin ו-Unpin, אופרנד buffer יכול להופיע בזוג אחד לכל היותר של כינויים, ופלט buffer חייב להופיע בזוג אחד של כינויים.

דוגמאות

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = {bar = 42 : i32},
  api_version = 4 : i32,
  called_computations <= [>@fo>o]
} : <(te>nsorf64) - tensorf64

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases< = [
    #stablehlo.output_operand_aliasoutput_tuple_indices = [],
      operand_ind>ex = 0,
     < oper>and>_tuple_<indic>es = []]
} : (memref4xf64) - memref4xf64

חילוק

סמנטיקה

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

• למספרים שלמים: חלוקת מספרים שלמים שמניבה את המנה האלגברית בלי החלק השברי.
• למספרים ממשיים: division מ-IEEE-754.
• למספרים מרוכבים: חילוק מרוכב.
• לסוגים שעברו קוונטיזציה:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

קלט

פלט

מגבלות

• (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(result).

דוגמאות

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs)< : (t>ensor4xf<32, t>ens>or4xf32<) - t>ensor4xf32
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

 דוגמאות נוספות

dot_general

סמנטיקה

מחשב מכפלה סקלרית בין פרוסות של lhs לבין פרוסות של rhs ומפיק טנזור result.

באופן רשמי יותר, result[result_index] = dot_product, כאשר:

• lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions].
• rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions].
• ‫result_batching_index + result_lhs_index + result_rhs_index = result_index כאשר size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) ו- size(result_rhs_index) = size(rhs_result_dimensions).
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions).
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :]).
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions)).
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions).
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :]).
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions)).
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y)).

לסוגים כמותיים, הפונקציה מבצעת dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result)).

עבור סוגים היברידיים של קוונטיזציה, הפעולה היא hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs).

‫precision_config שולט בפשרה בין מהירות לדיוק בחישובים במערכות עורפיות של מאיצים. אפשר להשתמש באחד מהערכים הבאים (בשלב הזה, הסמנטיקה של ערכי ה-enum האלה לא מוגדרת באופן מלא, אבל אנחנו מתכננים לטפל בזה ב-#755):

• ‫DEFAULT: החישוב הכי מהיר, אבל הקירוב הכי פחות מדויק למספר המקורי.
• HIGH: חישוב איטי יותר, אבל קירוב מדויק יותר למספר המקורי.
• ‫HIGHEST: החישוב הכי איטי, אבל הקירוב הכי מדויק למספר המקורי.

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

השדות DotAlgorithm כוללים:

• ‫lhs_precision_type ו-rhs_precision_type, רמות הדיוק שאליהן מעוגלים הצד הימני והצד השמאלי של הפעולה. סוגי הדיוק לא תלויים בסוגי האחסון של נתוני הקלט והפלט.
• accumulation_type רמת הדיוק שמשמשת לחישוב המצטבר.
• ‫lhs_component_count, rhs_component_count ו-num_primitive_operations רלוונטיים כשאנחנו מפעילים אלגוריתם שמפרק את הצד הימני או השמאלי של המשוואה למספר רכיבים ומבצע מספר פעולות מכפלה סקלרית 'פרימיטיביות' על הערכים האלה – בדרך כלל כדי לדמות דיוק גבוה יותר (לדוגמה, שימוש בסוג הנתונים של בינה מלאכותית bfloat16 לחישובים ברמת דיוק גבוהה יותר: bf16_6x tf32_3x וכו'). עבור אלגוריתמים ללא פירוק, הערכים האלה צריכים להיות 1.
• ‫allow_imprecise_accumulation כדי לציין אם מותר לצבור נתונים ברמת דיוק נמוכה יותר בחלק מהשלבים (למשל CUBLASLT_MATMUL_DESC_FAST_ACCUM).

מאפיינים DotAlgorithm לדוגמה:

// Inputs are casted to tf32, and then accumulated in f32:
{lhs_precision_type = tf32,
 rhs_precision_type = tf32,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = false}


// bf16_6x: each input is decomposed to 3 bf16 components, then 6 dot operations are done on those components, and the result is accumulated in f32.
{lhs_precision_type = bf16,
 rhs_precision_type = bf16,
 accumulation_type = f32,
 lhs_component_count = 3,
 rhs_component_count = 3,
 num_primitive_operations = 6,
 allow_imprecise_accumulation = false}


// Inputs are (casted to) f8e5m2, and we accumulate in f32, but for some steps we may accumulate in lower precision.
{lhs_precision_type = f8e5m2,
 rhs_precision_type = f8e5m2,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = true}

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

במאמר xla_data.proto > Algorithm מפורטים כמה ערכים נתמכים של אלגוריתמים. בכרטיס מספר 2483 מתואר התוכנית ליצירת מסמך מרכזי בנושא אלגוריתמים נתמכים לפי קצה עורפי.

קלט

פלט

מגבלות

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions).
• ‫(C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions).
• ‫(C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions).
• ‫(C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions).
• ‫(C5) 0 <= lhs_batching_dimensions < rank(lhs).
• ‫(C6) 0 <= lhs_contracting_dimensions < rank(lhs).
• ‫(C7) 0 <= rhs_batching_dimensions < rank(rhs).
• ‫(C8) 0 <= rhs_contracting_dimensions < rank(rhs).
• ‫(C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).
• ‫(C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...).
• ‫(C11) size(precision_config) = 2.
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions).
• אם הפעולה משתמשת בטנסורים לא מכומתים:
  • ‫(C13) element_type(lhs) = element_type(rhs).
• אם הפעולה משתמשת בטנסורים שעברו קוונטיזציה:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • ‫(C15) zero_points(rhs) = 0.
  • ‫(C16) אם is_per_axis_quantized(rhs), אז quantization_dimension(rhs) לא ב-rhs_contracting_dimensions.
  • אם is_quantized(lhs):
  • ‫(C17) storage_type(lhs) = storage_type(rhs).
  • ‫(C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • ‫(C19) אם is_per_tensor_quantized(rhs), אז is_per_tensor_quantized(result).
  • אם !is_quantized(lhs):
  • ‫(C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• אם !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation):
  • ‫(C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count.
  • ‫(C23) 0 < rhs_component_count.
  • ‫(C24) 0 < num_primitive_operations.

דוגמאות

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #sta<blehlo.dot
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimension>s = [1]
  ,
  precision_config = [<#stablehloprecisi>on DEFAULT, <#stablehloprecisi>on DEFAULT],
  algorithm = #stablehlo.dot<_algorithm
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation >= false
  
}< : (tenso>r2x2x2xi<64, tenso>r2x>2x2xi64<) - tenso>r2x2x2xi64
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

 דוגמאות נוספות

dynamic_broadcast_in_dim

סמנטיקה

הפעולה הזו זהה מבחינת הפונקציונליות ל-broadcast_in_dim op, אבל צורת התוצאה מוגדרת באופן דינמי באמצעות output_dimensions.

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

קלט

פלט

מגבלות

• ‫(C1) element_type(result) מחושב לפי הנוסחה הבאה:
  • ‫element_type(operand), אם !is_per_axis_quantized(operand).
  • ‫element_type(operand), למעט quantization_dimension(operand),‏ scales(operand) ו-zero_points(operand) שעשויים להיות שונים מ-quantization_dimension(result),‏ scales(result) ו-zero_points(result) בהתאמה, אחרת.
• ‫(C2) size(broadcast_dimensions) = rank(operand).
• ‫(C3) 0 <= broadcast_dimensions < rank(result).
• ‫(C4) is_unique(broadcast_dimensions).
• ‫(C5) לכל d בaxes(operand):
  • dim(operand, d) = 1 או
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• ‫(C6) אם is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • אם dim(operand, quantization_dimension(operand)) = 1, אז scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).
• ‫(C7) size(output_dimensions) = rank(result).
• ‫(C8) is_unique(known_expanding_dimensions + known_nonexpanding_dimensions).
• ‫(C9) 0 <= known_expanding_dimensions < rank(operand).
• ‫(C10) 0 <= known_nonexpanding_dimensions < rank(operand).

דוגמאות

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensio<ns = arra>yi64: 2, 1,
  known_expanding_dimensio<ns = a>rrayi64: 0,
  known_nonexpanding_dimensio<ns = a>rrayi64: 1
}< : (ten>sor1x3xi<64, t>ens>or3xi64<) - tenso>r2x3x2xi64
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 דוגמאות נוספות

dynamic_conv

סמנטיקה

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

קלט