מפרט StableHLO

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

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

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

כדי לצפות במפרט מגרסה קודמת של 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
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> הוא סוג t e n s o r f l o w, עם הצורה 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 הם אופציונליים ב-grammar, אבל ערכי ברירת המחדל שלהם הם 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, כולל הסוג, הערכים והשאלה האם יכולה להיות רק נקודה אחת או כמה נקודות אפסיות בסוג img_tensor. בהתאם לתוצאות הדיון, המפרט לגבי אפס נקודות עשוי להשתנות בעתיד (#1405).

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

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

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

בטנסורים שמבוצעת בהם כימות, אפשר לבצע כימות לכל טנסור, כלומר להשתמש ב-scale וב-zero_point אחד לכל הטנסור, או לבצע כימות לכל ציר, כלומר להשתמש במספר ערכים של 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}]

סוגי צמדי-ערכי-מפתח (Tuple) מייצגים צמדי-ערכי-מפתח, כלומר רשימות הטרוגניות. צמדי מפתח/ערך הם תכונה מדור קודם שקיימת רק לצורך תאימות ל-HLO. ב-HLO, צמדים משמשים לייצוג משתני קלט ופלט שונים. ב-StableHLO יש תמיכה בקלט ופלט שונים, והשימוש היחיד בפלטים ב-StableHLO הוא לייצג באופן מקיף HLO ABI כאשר למשל T, tuple<T> ו-tuple<tuple<T>> עשויים להיות שונים באופן מהותי בהתאם להטמעה מסוימת. בעתיד, אנחנו מתכננים לבצע שינויים ב-HLO ABI, שיכולים לאפשר לנו להסיר סוגים שונים של Tuple מ-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 באמצעות ערכים של טינסורים ב-0 מימדים מסוג 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 ביט לרשתות נוירונליות עמוקות.
  • סוג f8E4M3B11FNUZ התואם לקידוד E4M3 של הפורמטים של FP8 שמתוארים במאמר Hybrid 8-bit Floating Point (HFP8) Training and Inference for Deep Neural Networks.
  • סוג bf16 התואם לפורמט bfloat16 שמתואר במאמר BFloat16: הסוד לביצועים גבוהים ב-Cloud TPUs.
  • הסוגים f16, ‏ f32 ו-f64 תואמים לפורמטים binary16 ('חצי דיוק'), ‏ binary32 ('דיוק יחיד') ו-binary64 ('דיוק כפול') שמתוארים בתקן IEEE 754.
  • הסוג tf32 תואם לפורמט TensorFloat32, ויש לו תמיכה מוגבלת ב-StableHLO.
  • סוגי MX (מיקרו-שינוי) מסוגים f4E2M1FN,‏ f6E2M3FN,‏ f6E3M2FN ו-f8E8M0FNU שמתוארים במפרט הפורמטים של OCP Microscaling.
• סוגי רכיבים מורכבים מייצגים ערכים מורכבים שיש להם חלק ממשי וחלק מדומה מאותו סוג רכיב. הסוגים המורכבים הנתמכים הם 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 – ליצור יכולת פעולה הדדית בין מסגרות למידת מכונה לבין מהדרי למידת מכונה.

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

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

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

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

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

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

התחביר של פונקציות הקלט כולל חלק שלא בשימוש כרגע (ראו את הפונקציה Unused למעלה), שנועד לצורך תאימות ל-MLIR. ב-MLIR יש מושג כללי יותר של 'אזורים', שיכולים לכלול כמה 'בלוקים' של פעולות שמחוברים יחד באמצעות פעולות מעבר (jump). למקטעים האלה יש מזהי 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> מייצג ערך t e n s o r f l o w, עם המיפוי הבא מאינדקסים לאלמנטים: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. הסדר שבו הרכיבים האלה מאוחסנים בזיכרון מוגדר בהטמעה. למשתני הטנזור הקבועים יש את האילוצים הבאים:

• (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) '>'

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

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

תפעול

בטן

סמנטיקה

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

• למספרים שלמים בעלי סימן: מודולוס של מספר שלם.
• למספרים מסוג float: 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. בהתאם לסוג הרכיב, הפעולות הבאות:

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

קלט

פלט

אילוצים

• אם הפעולה משתמשת בטנסורים לא מקובצים:
  • (C1) type(lhs) = type(rhs) = type(result).
• אם הפעולה משתמשת בטנסורים מרובים (quantization):
  • (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.

הפעולה מפצלת את רשת התהליכים של 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 של שני טינסורים 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) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[1, 2], [3, 0]]

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

atan2

סמנטיקה

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

• למספרים מסוג float: 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) : (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. באופן רשמי יותר, אפשר לבטא את הפעולה הזו מפירוק לפעולות 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, ומנורמלת את הטנזור 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
} : (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

סמנטיקה

הפונקציה מבצעת פעולת 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) : (tensor<f64>) -> tensor<4xf16>
// %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 אם 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_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

סמנטיקה

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

• למספרים מסוג float: 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

סמנטיקה

הפונקציה מבצעת חישוב של חזקה שלם של רכיבי הטנזור 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) : (tensor<5xf32>) -> tensor<5xf32>
// %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 שבו מטריצת הקלט היא לא מטריצת Hermitian חיובית-מוגדרת, ההתנהגות לא מוגדרת.

לסוגי נתונים שמפורטים במספרים, הפונקציה מבצעת את הפעולה 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.

הפעולה מפצלת את רשת התהליכים של 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, שולחים את הערך של המכפלה operand מתהליך המקור לתהליך היעד ויוצרים מכפלה result.

הפעולה מפצלת את רשת התהליכים של 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]]

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

compare

סמנטיקה

הפונקציה מבצעת השוואה של הרכיבים של הטנסורים 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 אם 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]]

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

להשלים המרה

סמנטיקה

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

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

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

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

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

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

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

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

convolve

סמנטיקה

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

• למספרים מסוג float: 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 ומפיקה את טנזור 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 כדי לספק מטא-נתונים נוספים שמוגדרים בהטמעה.

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

• למספרים שלמים: חלוקת מספרים שלמים שמניבה את החלק האלגברי, בלי החלק העשרוני.
• למספרים מסוג float: 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 קובע את האיזון בין המהירות לדיוק של החישובים בקצוות העורפי של המאיצים. הערך יכול להיות אחד מהערכים הבאים (כרגע, הסמנטיקה של ערכי המניין האלה לא מוגדרת במלואה, אבל אנחנו מתכננים לטפל בבעיה הזו בבקשה מספר 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).
• אם הפעולה משתמשת בטנסורים מרובים (quantization):
  • (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_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_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_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

סמנטיקה

הפעולה הזו זהה מבחינה פונקציונלית לפעולת convolution, אבל המילוי מצוין באופן דינמי באמצעות 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]]
//          ]]