מפרט StableHLO

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

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

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

כדי להציג את המפרט מגרסה קודמת של StableHLO, פותחים את המאגר בכתובת גרסה מתויגת של תחום עניין. לדוגמה, המפרט של StableHLO v0.19.0. כדי לראות את השינויים שהתרחשו בכל העברה משנית של 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
NonValueType ::= TensorElementType | QuantizedTensorElementType | FunctionType | StringType

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

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

סוגים של חיישנים מייצגים טנסטורים, כלומר מערכים רב-מימדיים. יש להם הצורה וסוג רכיב, שבהם צורה מייצגת מספר לא שלילי או גדלי מאפיינים לא ידועים בסדר עולה לפי הסדר מאפיינים (שנקראים גם צירים) ממוספרים מ-0 עד R-1. מספר המאפיינים R נקרא דירוג. לדוגמה, tensor<2x3xf32> הוא סוג img_tensor עם הצורה 2x3 וסוג הרכיב f32. יש לו שני מאפיינים (או, במילים אחרות, שני צירים) - מימד 0 ומימד ראשון - שגודלו הם 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 ::= IntegerConstant
QuantizationStorageMax ::= IntegerConstant
QuantizationExpressedType ::= FloatType
QuantizationDimension ::= IntegerConstant
QuantizationParameters ::= QuantizationParameter
                         | '{' QuantizationParameter {',' QuantizationParameter} '}'
QuantizationParameter ::= QuantizationScale ':' QuantizationZeroPoint
QuantizationScale ::= FloatConstant
QuantizationZeroPoint ::= IntegerConstant

סוגי רכיבים בכמות גדולה מייצגים ערכים שלמים של סוג האחסון ב- הטווח בין storage_min ל-storage_max (כולל) שתואם ל- ערכי נקודה צפה (floating-point) מסוג מבטא. במספר שלם נתון 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 הוא קבוע מסוג נקודה צפה (floating-point), אבל יש עניין רב בסולמות המבוססים על מספרים שלמים, המיוצגים באמצעות מכפילים אנחנו מתכננים לבדוק את זה בעתיד הקרוב (#1404).

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

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

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

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

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

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

TokenType ::= 'token'

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

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

סוגי צמדים מייצגים צמדים, כלומר רשימות הטרוגניות. קפלים הם מורשת לתאימות ל-HLO בלבד. ב-HLO, צמדים משמש לייצוג קלט ופלט משתנים. ב-StableHLO, קלט ווריאציה יש תמיכה בפלט באופן מקורי, והשימוש היחיד ב-Tuples ב-StableHLO הוא לייצג באופן מקיף את ממשק ה-ABI של HLO, T, tuple<T> וגם הערך של tuple<tuple<T>> עשוי להיות שונה באופן מהותי בהתאם לישות מסוימת יישום בפועל. בעתיד, אנחנו מתכננים לבצע שינויים ב-HLO ABI ועשויה להסיר מ-StableHLO סוגים שונים של tuple (#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 ::= 'f8E4M3FN' | 'f8E5M2' | 'f8E4M3FNUZ' | 'f8E5M2FNUZ'
            | 'f8E4M3B11FNUZ' | '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, כולל.
• סוגים של נקודות צפות יכולים להיות אחד מהבאים:
  • הסוגים של f8E4M3FN ו-f8E5M2 שתואמים לסוג הקידודים E4M3 ו-E5M2 של פורמט FP8 שמתוארים ב: פורמטי FP8 ללמידה עמוקה (Deep Learning).
  • סוגי f8E4M3FNUZ וf8E5M2FNUZ שתואמים לE4M3 ולE5M2 של פורמטי FP8 שמתוארים פורמטים מספריים של 8 ביט לרשתות נוירונים עמוקות
  • סוג f8E4M3B11FNUZ שתואם לקידוד E4M3 של פורמטי FP8 מתואר ב: אימון של נקודה צפה (HFP8) היברידית והסקת מסקנות לגבי רשתות נוירונים עמוקות
  • סוג bf16 שתואם לפורמט bfloat16 שמתואר ב: BFloat16: הסוד לביצועים גבוהים במעבדי Cloud TPU.
  • הסוגים של f16, f32 ו-f64 תואמים בהתאמה binary16 ('half precision'), binary32 ('דיוק יחיד') וגם הפורמטים binary64 (דיוק כפול) שמתוארים בפורמט תקן IEEE 754.
  • סוג tf32 תואם לפורמט TensorFloat32 והיא מספקת תמיכה מוגבלת ב-StableHLO.
• סוגים מורכבים מייצגים ערכים מורכבים שיש להם חלק ממשי וחלק דמיוני מאותו סוג רכיב. מתחם נתמך הסוגים הם complex<f32> (שני החלקים הם מסוג f32) ו-complex<f64> (שני החלקים הם מסוג f64).

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

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

StringType ::= 'string'

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

תפעול

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

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 פונקציות הן לא ערכים מהמחלקה הראשונה, במאפייני הקלט (שמופיעים גם בצורה סטטית). סוגי הקלט והפלט צריכה ומופקת על ידי אופ תלוי במנמן שלו. לדוגמה, add הפונקציה op צריכה 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 בטבע כוללות לפעמים מאפיינים שלא מתוארות במסמך הזה. בעתיד, אנחנו מתכננים לקלוט את המאפיינים האלה בפונקציה 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]

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

• (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>. מייצג ערך t e n s o r f l o עם המיפוי הבא מאינדקסים לאלמנטים: {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) '>'

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

• (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.

תפעול

בטן

סמנטיקה

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

• למספרים שלמים חתומים: מודול של מספרים שלמים.
• לצפים: 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)) אם is_complex(operand).
  • אחרת, baseline_element_type(operand).

דוגמאות

// %operand: [-2, 0, 2]
%result = "stablehlo.abs"(%operand) : (tensor<3xi32>) -> tensor<3xi32>
// %result: [2, 0, 2]

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

add

סמנטיקה

מבצעת הוספה מבוססת-רכיבים של שני טנזורים 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) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %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, !stablehlo.token) -> !stablehlo.token

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

all_gather

סמנטיקה

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

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

• cross_replica(replica_groups) אם channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups) אם channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups) אם 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...) מלבד:
  • 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_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %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) אם channel_id <= 0 and use_global_device_ids = false.
• cross_replica_and_partition(replica_groups) אם channel_id > 0 and use_global_device_ids = false.
• flattened_ids(replica_groups) אם 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(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %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) אם channel_id <= 0.
• cross_partition(replica_groups) אם 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_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %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 (וגם) מבחינת רכיבי AND של שני טנסורים lhs ו-rhs ומפיק result את Tensor. בהתאם לסוג הרכיב, הפעולות הבאות:

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

קלט

פלט

מגבלות

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

דוגמאות

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

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

atan2

סמנטיקה

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

• לצפים: atan2 מ-IEEE-754.
• למספרים מרוכבים: 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) : (tensor<3xf64>, tensor<3xf64>) -> tensor<3xf64>
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

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

batch_norm_grad

סמנטיקה

מחשב הדרגתיות של מספר קלטים של הדבקה לאחור של batch_norm_training מ-grad_output, ומפיקה את grad_operand, grad_scale ו-grad_offset את Tensors. באופן רשמי יותר, ניתן לבטא פעולה זו כפירוק פעולות 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
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>,
     tensor<2x2x2xf64>) -> (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %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
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) -> tensor<2x2x2xf64>
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

סמנטיקה

הפונקציה מחשבת את הממוצע והשונות בכל המאפיינים מלבד feature_index ומנורמל את ערך ה-tensor operand, שמפיק output, batch_mean ו-batch_var tensor. יותר רשמי, אפשר לבטא את הפעולה הזו פירוק לפעולות 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
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>) ->
    (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %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

סמנטיקה

מבצע פעולת ביט-cast על טנזור 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 מחזירה ייצוג בזיכרון של ערך נתון, ואת ההתנהגות שלו מוגדר ביישום כי הייצוג המדויק של tensor והייצוג המדויק של סוגי הרכיבים, מוגדר גם כן.

קלט

פלט

מגבלות

• (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) : (tensor<f64>) -> tensor<4xf16>
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

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

broadcast_in_dim

סמנטיקה

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

• operand_index[d] = 0 אם 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) resp, אחרת.
• (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_dimensions = array<i64: 2, 1>
} : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

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

כיסוי

סמנטיקה

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

• selected_branch = branches[index] אם 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, %result_branch0) : (tensor<2xi64>, tensor<2xi64>) -> ()
}, {
  "stablehlo.return"(%result_branch1, %result_branch1) : (tensor<2xi64>, tensor<2xi64>) -> ()
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
// %result0: [1, 1]
// %result1: [1, 1]

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

cbrt

סמנטיקה

מבצעת פעולת שורש ממעלה שלישית ברמת הרכיב ב-Tensor operand ומפיקה result t e n s o r f l o w. בהתאם לסוג הרכיב, הפעולות הבאות:

• לצפים: 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) : (tensor<4xf64>) -> tensor<4xf64>
// %result: [0.0, 1.0, 2.0, 3.0]

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

ceil

סמנטיקה

מבצע ceil מקודד ברמת היסוד של operand tenor ומפיק 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) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

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

צ'ולסקי

סמנטיקה

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

באופן רשמי יותר, לכל ה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
} : (tensor<3x3xf32>) -> tensor<3x3xf64>
// %result: [
//           [1.0, 0.0, 0.0],
//           [2.0, 4.0, 0.0],
//           [3.0, 5.0, 6.0]
//          ]

הצמדה

סמנטיקה

מצמידה כל רכיב של הטנזור operand בין מינימום למקסימום ומפיק טנזור result. יותר רשמי: 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) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
// %result: [5, 13, 20]

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

collective_broadcast

סמנטיקה

בכל קבוצת תהליכים ברשת התהליכים של StableHLO, שולחים את הערך של טנזור operand מתהליך המקור לתהליכי היעד ולהפיק result t e n s o r f l o w.

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

• cross_replica(replica_groups) אם channel_id <= 0.
• cross_partition(replica_groups) אם 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_groups = dense<[[2, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

סמנטיקה

בתוך כל קבוצת תהליכים ברשת התהליכים של StableHLO, שולח הערך של Tensor של operand מתהליך המקור לתהליך היעד ומפיק result t e n s o r f l o w.

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

• cross_replica(source_target_pairs) אם channel_id <= 0.
• cross_partition(source_target_pairs) אם 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_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x2xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[0, 0], [0, 0]]
// %result@(1, 0): [[1, 2], [3, 4]]
// %result@(2, 0): [[5, 6], [7, 8]]

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

השוואה

סמנטיקה

מבצע השוואה מבחינת הרכיבים של lhs ו-rhs Tensors בהתאם 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.

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

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

לסוגי אלמנטים של נקודה צפה (floating-point) עם 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 אם is_signed_integer(element_type(lhs)).
  • UNSIGNED אם is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT או TOTALORDER אם is_float(element_type(lhs)).
  • FLOAT אם is_complex(element_type(lhs)).

דוגמאות

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %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) : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xcomplex<f64>>
// %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,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

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

לשרשר

סמנטיקה

הפונקציה משרשרת את 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]) מלבד:
  • 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
} : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

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

קבוע

סמנטיקה

מפיקה טנזור output מקבוע value.

קלט

פלט

מגבלות

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

דוגמאות

%output = "stablehlo.constant"() {
  value = dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
} : () -> tensor<2x2xf32>
// %output: [[0.0, 1.0], [2.0, 3.0]]

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

להשלים המרה

סמנטיקה

מבצע המרה ברמת הרכיבים מסוג רכיב אחד לסוג אחר ב- t e n s o r f l o w, ומפיק טנזור result.operand

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

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

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

המרות מסוג מורכב למורכבות פועלות באותה התנהגות המרות floating-point-to-floating-point שמיועדות להמרות ריאליות חלקים דמיוניים.

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

בעיקרון, פעולה זו יכולה לבטא את ביטול החלוקה (המרה מ- טנסטורים כמותיים לטינורים רגילים), קוונטיזציה (המרה מהגדרות רגילות) img_tensors to quantor tensors (טנזטורים כמותיים) ו-requantization (המרה בין ערכים כמותיים) את Tensors), אבל כרגע יש לנו פעולות ייעודיות לכך, uniform_dequantize לתרחיש הראשון ו-uniform_quantize לתרחיש לדוגמה והשלב השני והשלישי. בעתיד, ייתכן ששתי האפשרויות האלה ימוזגו אל convert (#1576).

קלט

פלט

מגבלות

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

דוגמאות

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand) : (tensor<3xi64>) -> tensor<3xcomplex<f64>>
// %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 אם result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) אם 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_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

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

קוסינוס

סמנטיקה

מבצעת פעולת קוסינוס מבוססת-רכיבים על טנזור operand ומפיקה result t e n s o r f l o w. בהתאם לסוג הרכיב, הפעולות הבאות:

• לצפים: 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) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.0], [-1.0, 0.0]]

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

count_leading_zeros

סמנטיקה

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

קלט

פלט

מגבלות

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

דוגמאות

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

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

custom_call

סמנטיקה

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

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

קלט

פלט

דוגמאות

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

חילוק

סמנטיקה

מבצע חלוקה חכמה של הרכיבים של 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) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// %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, הדיוק של ה-LHS RHS של הפעולה מעוגל ל. סוגי הדיוק אינם תלויים של סוגי הקלט והפלט.
• accumulation_type רמת הדיוק שמשמשת לצבירה.
• lhs_component_count, rhs_component_count וגם num_primitive_operations כאשר אנחנו מבצעים אלגוריתם שמפרק את ה-LHS ו/או ה-RHS ויוצרת ריבוי רכיבים פעולות בנקודות - בדרך כלל כדי לבצע אמולציה של רמת דיוק גבוהה יותר (למשל ניצול סוג הנתונים bfloat16 של בינה מלאכותית (AI) לביצוע חישובים מדויקים יותר: 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 = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision 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
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

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

dynamic_broadcast_in_dim

סמנטיקה

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

הפעולה הזו מקבלת גם את המאפיינים האופציונליים known_expanding_dimensions, known_non_expanding_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) resp, אחרת.
• (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_non_expanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_non_expanding_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_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_non_expanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

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

dynamic_conv

סמנטיקה

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

קלט

פלט

מגבלות

• (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 אם result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) אם 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]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]