مشخصات StableHLO

StableHLO یک مجموعه عملیات برای عملیات سطح بالا (HLO) در مدل‌های یادگیری ماشین (ML) است. StableHLO به عنوان یک لایه قابل حمل بین چارچوب‌های مختلف یادگیری ماشین و کامپایلرهای یادگیری ماشین عمل می‌کند: چارچوب‌های یادگیری ماشینی که برنامه‌های StableHLO تولید می‌کنند با کامپایلرهای یادگیری ماشینی که برنامه‌های StableHLO را مصرف می‌کنند، سازگار هستند.

هدف ما ساده‌سازی و تسریع توسعه‌ی یادگیری ماشین با ایجاد قابلیت همکاری بیشتر بین چارچوب‌های مختلف یادگیری ماشین (مانند TensorFlow، JAX و PyTorch) و کامپایلرهای یادگیری ماشین (مانند XLA و IREE) است. برای این منظور، این سند مشخصات زبان برنامه‌نویسی StableHLO را ارائه می‌دهد.

این مشخصات شامل سه بخش اصلی است. اول، بخش برنامه‌ها ساختار برنامه‌های StableHLO را شرح می‌دهد که شامل توابع StableHLO هستند که خود شامل عملیات StableHLO هستند. در این ساختار، بخش عملیات، معنای عملیات‌های منفرد را مشخص می‌کند. بخش اجرا ، معنای تمام این عملیات‌ها را که با هم در یک برنامه اجرا می‌شوند، ارائه می‌دهد. در نهایت، بخش نمادگذاری، نمادگذاری مورد استفاده در سراسر مشخصات را مورد بحث قرار می‌دهد.

برای مشاهده مشخصات نسخه‌های قبلی StableHLO، مخزن مربوط به نسخه مورد نظر را باز کنید. برای مثال، مشخصات StableHLO نسخه ۰.۱۹.۰ . برای مشاهده تغییراتی که در هر نسخه جزئی StableHLO رخ داده است، به گزارش نسخه در VhloDialect.td مراجعه کنید.

برنامه‌ها

Program ::= {Func}

برنامه‌های StableHLO شامل تعداد دلخواهی از توابع StableHLO هستند. در زیر یک برنامه نمونه با تابع @main شده است که دارای ۳ ورودی ( %image ، %weights و %bias ) و ۱ خروجی است. بدنه تابع دارای ۶ عملیات است.

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 مشابه شناسه‌ها در بسیاری از زبان‌های برنامه‌نویسی هستند، با دو ویژگی: ۱) همه شناسه‌ها دارای علامت‌هایی هستند که انواع مختلف شناسه‌ها را از هم متمایز می‌کنند، ۲) شناسه‌های مقدار می‌توانند کاملاً عددی باشند تا تولید برنامه‌های StableHLO ساده شود.

انواع

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

انواع StableHLO به انواع مقداری (که انواع درجه یک نیز نامیده می‌شوند) که نشان‌دهنده مقادیر StableHLO هستند و انواع غیر مقداری که سایر عناصر برنامه را توصیف می‌کنند، طبقه‌بندی می‌شوند. انواع StableHLO مشابه انواع در بسیاری از زبان‌های برنامه‌نویسی هستند، با این تفاوت که ویژگی اصلی آن ماهیت وابسته به دامنه StableHLO است که منجر به برخی نتایج غیرمعمول می‌شود (مثلاً انواع اسکالر، انواع مقداری نیستند).

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

انواع تانسور ، تانسورها، یعنی آرایه‌های چندبعدی را نشان می‌دهند. آن‌ها دارای یک شکل و یک نوع عنصر هستند که در آن یک شکل ، اندازه‌های ابعاد غیر منفی یا ناشناخته را به ترتیب صعودی ابعاد مربوطه (که به آن‌ها محور نیز گفته می‌شود) نشان می‌دهد که از 0 تا R-1 شماره‌گذاری شده‌اند. تعداد ابعاد R رتبه نامیده می‌شود. به عنوان مثال، tensor<2x3xf32> یک نوع تانسور با شکل 2x3 و نوع عنصر f32 است. این نوع تانسور دو بعد (یا به عبارت دیگر، دو محور) دارد - بعد 0 و بعد 1 - که اندازه‌های آن‌ها 2 و 3 است. رتبه آن 2 است.

شکل‌ها می‌توانند تا حدی یا کاملاً ناشناخته باشند (پویا)، مثلاً tensor<?x2xf64> تا حدی ناشناخته و tensor<?x?xf64> کاملاً ناشناخته است. اندازه‌های ابعاد پویا با استفاده از ? نمایش داده می‌شوند. شکل‌ها را نمی‌توان رتبه‌بندی کرد.

در آینده، ما قصد داریم گسترش انواع تانسور را فراتر از اندازه ابعاد و انواع عناصر بررسی کنیم، به عنوان مثال، شامل طرح‌بندی‌ها ( #629 ) و پراکندگی ( #1078 ).

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

انواع عناصر کوانتیزه، مقادیر صحیح از یک نوع ذخیره‌سازی را در محدوده storage_min تا storage_max (شامل) نشان می‌دهند که با مقادیر ممیز شناور از یک نوع بیان‌شده مطابقت دارند. برای یک مقدار صحیح داده شده i ، مقدار ممیز شناور مربوطه f می‌توان به صورت f = (i - zero_point) * scale محاسبه کرد، که در آن scale و zero_point پارامترهای کوانتیزه نامیده می‌شوند. storage_min و storage_max در گرامر اختیاری هستند، اما به ترتیب مقادیر پیش‌فرض min_value(storage_type) و max_value(storage_type) را دارند. انواع عناصر کوانتیزه محدودیت‌های زیر را دارند:

• (C1) type(storage_min) = storage_type .
• (C2) type(storage_max) = storage_type .
• (C3) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type)
• type(scales...) = expressed_type .
• (C5) 0 < scales .
• (C6) is_finite(scales...) .
• (C7) storage_min <= zero_points <= storage_max .
• (C8) type(zero_points...) = storage_type .
• (C9) size(scales) = size(zero_points) .
• (C10) اگر is_empty(quantization_dimension) باشد، آنگاه size(scales) = 1 .
• (C11) 0 <= quantization_dimension .

در حال حاضر، QuantizationScale یک ثابت ممیز شناور است، اما علاقه زیادی به مقیاس‌های مبتنی بر عدد صحیح، که با ضرایب و شیفت‌ها نمایش داده می‌شوند، وجود دارد. ما قصد داریم این موضوع را در آینده نزدیک ( #1404 ) بررسی کنیم.

بحثی مداوم در مورد معناشناسی QuantizationZeroPoint ، از جمله نوع، مقادیر و اینکه آیا می‌تواند فقط یک یا چندین نقطه صفر در یک نوع تانسور کوانتیزه وجود داشته باشد، وجود دارد. بر اساس نتایج این بحث، مشخصات مربوط به نقاط صفر ممکن است در آینده تغییر کند ( #1405 ).

بحث دیگری که در حال انجام است، شامل معناشناسی QuantizationStorageMin و QuantizationStorageMax است تا مشخص شود که آیا باید محدودیتی بر این مقادیر و بر مقادیر تانسورهای کوانتیزه شده اعمال شود یا خیر ( #1406 ).

در نهایت، ما قصد داریم نمایش مقیاس‌های ناشناخته و نقاط صفر را بررسی کنیم، مشابه روشی که قصد داریم نمایش اندازه‌های ابعاد ناشناخته ( #1407 ) را بررسی کنیم.

انواع تانسور کوانتیزه، تانسورهایی با عناصر کوانتیزه را نشان می‌دهند. این تانسورها دقیقاً مشابه تانسورهای معمولی هستند، با این تفاوت که عناصر آنها به جای انواع عناصر معمولی، دارای انواع عناصر کوانتیزه هستند.

در تانسورهای کوانتیزه، کوانتیزاسیون می‌تواند per-tensor باشد، به این معنی که برای کل تانسور یک scale و zero_point داشته باشد یا می‌تواند per-axis باشد، به این معنی که چندین scales و zero_points داشته باشد، به این معنی که برای هر برش از یک بُعد خاص، یک جفت quantization_dimension وجود داشته باشد. به طور رسمی‌تر، در یک تانسور t با کوانتیزاسیون per-axis، برش‌های dim(t, quantization_dimension) از بُعد quantization_dimension وجود دارد: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] و غیره. همه عناصر در برش i ام از scales[i] و zero_points[i] به عنوان پارامترهای کوانتیزاسیون خود استفاده می‌کنند. انواع تانسور کوانتیزه محدودیت‌های زیر را دارند:

• برای کوانتیزاسیون در هر تانسور:
  • بدون هیچ قید و بند اضافی.
• برای کوانتیزاسیون در هر محور:
  • (C12) quantization_dimension < rank(self) .
  • (C13) dim(self, quantization_dimension) = size(scales)

TokenType ::= 'token'

انواع توکن ، نشان‌دهنده‌ی توکن‌ها هستند، یعنی مقادیر مبهمی که توسط برخی عملیات تولید و مصرف می‌شوند. توکن‌ها برای اعمال ترتیب اجرا بر عملیات‌ها، همانطور که در بخش اجرا توضیح داده شده است، استفاده می‌شوند.

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

انواع بافر نشان‌دهنده بافرها هستند. برای مثال، در XLA، بافرها آرایه‌های چندبعدی با ذخیره‌سازی سازگار هستند. مشابه انواع تانسور ، انواع بافر دارای یک شکل و یک نوع عنصر هستند، که در آن یک شکل نشان‌دهنده اندازه‌های ابعاد غیر منفی یا ناشناخته به ترتیب صعودی ابعاد مربوطه (که به آنها محور نیز گفته می‌شود) است که از 0 تا R-1 شماره‌گذاری شده‌اند. تعداد ابعاد R رتبه نامیده می‌شود. برای مثال، memref<2x3xf32> یک نوع بافر با شکل 2x3 و نوع عنصر f32 است. این نوع بافر دو بعد (یا به عبارت دیگر، دو محور) دارد - بعد 0 و بعد 1 - که اندازه‌های آنها 2 و 3 است. رتبه آن 2 است.

بافرها را می‌توان با استفاده از custom_call به CreateBuffer یا Pin اختصاص داد و از طریق custom_call به Unpin ، تخصیص را آزاد کرد. فقط عملیات custom_call می‌توانند محتوای داخل بافرها را بخوانند و بنویسند. برای جزئیات بیشتر به custom_call مراجعه کنید.

انواع تاپل نشان‌دهنده‌ی تاپل‌ها، یعنی لیست‌های ناهمگن هستند. تاپل‌ها یک ویژگی قدیمی هستند که فقط برای سازگاری با HLO وجود دارند. در HLO، از تاپل‌ها برای نمایش ورودی‌ها و خروجی‌های متغیر استفاده می‌شود. در StableHLO، ورودی‌ها و خروجی‌های متغیر به صورت بومی پشتیبانی می‌شوند و تنها کاربرد تاپل‌ها در StableHLO نمایش جامع HLO ABI است که در آن مثلاً T ، tuple<T> و tuple<tuple<T>> ممکن است بسته به یک پیاده‌سازی خاص، از نظر مادی متفاوت باشند. در آینده، ما قصد داریم تغییراتی در HLO ABI ایجاد کنیم که ممکن است به ما امکان حذف انواع تاپل از StableHLO ( #598 ) را بدهد.

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

انواع عناصر، عناصری از انواع تانسور را نشان می‌دهند. برخلاف بسیاری از زبان‌های برنامه‌نویسی، این نوع‌ها در StableHLO درجه یک نیستند. این بدان معناست که برنامه‌های StableHLO نمی‌توانند مستقیماً مقادیر این نوع‌ها را نشان دهند (در نتیجه، اصطلاحاً گفته می‌شود که مقادیر اسکالر از نوع T را با مقادیر تانسور صفر بعدی از نوع tensor<T> نشان دهیم).

• نوع بولی (Boolean) نشان دهنده مقادیر بولی 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 که از قراردادهای IEEE-754 پیروی می‌کنند.
  • انواع f8E4M3FN و f8E5M2 که به ترتیب مربوط به کدگذاری‌های E4M3 و E5M2 از فرمت FP8 هستند و در بخش «فرمت‌های FP8 برای یادگیری عمیق» توضیح داده شده‌اند.
  • انواع f8E4M3FNUZ و f8E5M2FNUZ مربوط به کدگذاری‌های E4M3 و E5M2 از فرمت‌های FP8 که در «فرمت‌های عددی ۸ بیتی برای شبکه‌های عصبی عمیق» شرح داده شده‌اند.
  • نوع f8E4M3B11FNUZ مربوط به کدگذاری E4M3 از فرمت‌های FP8 که در آموزش و استنتاج ممیز شناور 8 بیتی ترکیبی (HFP8) برای شبکه‌های عصبی عمیق شرح داده شده است.
  • نوع bf16 مربوط به قالب bfloat16 که در کتاب BFloat16: راز عملکرد بالا در TPU های ابری توضیح داده شده است.
  • انواع f16 ، f32 و f64 که به ترتیب مربوط به قالب‌های binary16 ("دقت نیم") ، binary32 ("دقت تک") و binary64 ("دقت دو برابر") هستند که در استاندارد IEEE 754 شرح داده شده‌اند.
  • نوع tf32 با فرمت TensorFloat32 مطابقت دارد و پشتیبانی محدودی در StableHLO دارد.
  • انواع f4E2M1FN ، f6E2M3FN ، f6E3M2FN و f8E8M0FNU MX (ریزمقیاس‌بندی) که در مشخصات قالب‌های ریزمقیاس‌بندی OCP شرح داده شده‌اند.
• انواع مختلط، مقادیر مختلطی را نشان می‌دهند که دارای یک بخش حقیقی و یک بخش موهومی از یک نوع عنصر هستند. انواع مختلط پشتیبانی شده عبارتند از complex<f32> (هر دو بخش از نوع f32 هستند) و complex<f64> (هر دو بخش از نوع f64 هستند).

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

انواع تابع ، هم توابع با نام و هم توابع ناشناس را نشان می‌دهند. آن‌ها دارای انواع ورودی (لیست انواع در سمت چپ -> ) و انواع خروجی (لیست انواع در سمت راست -> ) هستند. در بسیاری از زبان‌های برنامه‌نویسی، انواع تابع درجه یک هستند، اما در StableHLO اینطور نیست.

StringType ::= 'string'

نوع رشته (String) توالی‌هایی از بایت‌ها را نشان می‌دهد. برخلاف بسیاری از زبان‌های برنامه‌نویسی، نوع رشته در StableHLO درجه یک نیست و فقط برای تعیین فراداده‌های استاتیک برای عناصر برنامه استفاده می‌شود.

عملیات

عملیات StableHLO (که به آنها ops نیز گفته می‌شود) مجموعه‌ای بسته از عملیات سطح بالا در مدل‌های یادگیری ماشین را نشان می‌دهند. همانطور که در بالا بحث شد، سینتکس StableHLO به شدت از MLIR الهام گرفته شده است، که لزوماً ارگونومیک‌ترین جایگزین نیست، اما مسلماً بهترین گزینه برای هدف StableHLO است که ایجاد قابلیت همکاری بیشتر بین چارچوب‌های یادگیری ماشین و کامپایلرهای یادگیری ماشین است.

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 دو مقدار ورودی مصرف کرده و یک مقدار خروجی تولید می‌کند. در مقایسه، عملیات select_and_scatter سه مقدار ورودی، دو تابع ورودی و سه ویژگی ورودی مصرف می‌کند.

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

توابع ورودی (که توابع ناشناس نیز نامیده می‌شوند) بسیار شبیه به توابع نامگذاری شده هستند، به جز اینکه: ۱) شناسه ندارند (از این رو "ناشناس" نامیده می‌شوند)، ۲) انواع خروجی را تعریف نمی‌کنند (انواع خروجی از مقدار return درون تابع استنباط می‌شوند).

نحو توابع ورودی شامل یک بخش بلااستفاده فعلی است (به بخش تولید Unused در بالا مراجعه کنید) که برای سازگاری با MLIR وجود دارد. در MLIR، مفهوم کلی‌تری از "مناطق" وجود دارد که می‌توانند چندین "بلوک" از عملیات را داشته باشند که از طریق عملیات پرش به یکدیگر متصل شده‌اند. این بلوک‌ها دارای شناسه‌هایی هستند که با تولید Unused مطابقت دارند، به طوری که می‌توان آنها را از یکدیگر تشخیص داد. StableHLO عملیات پرش ندارد، بنابراین بخش مربوطه از نحو MLIR بلااستفاده است (اما هنوز وجود دارد).

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

ویژگی‌های ورودی دارای یک نام و یک مقدار هستند که یکی از ثابت‌های پشتیبانی‌شده است. آن‌ها روش اصلی برای تعیین فراداده‌های استاتیک برای عناصر برنامه هستند. به عنوان مثال، تابع concatenate op از ویژگی dimension برای تعیین بُعدی که مقادیر ورودی آن در امتداد آن الحاق می‌شوند، استفاده می‌کند. به طور مشابه، تابع slice op از چندین ویژگی مانند start_indices و limit_indices برای تعیین مرزهایی که برای برش مقدار ورودی استفاده می‌شوند، استفاده می‌کند.

در حال حاضر، برنامه‌های StableHLO در حالت عادی گاهی اوقات حاوی ویژگی‌هایی هستند که در این سند توضیح داده نشده‌اند. در آینده، ما قصد داریم این ویژگی‌ها را در مجموعه StableHLO جذب کنیم یا از نمایش آنها در برنامه‌های StableHLO جلوگیری کنیم. در عین حال، لیست این ویژگی‌ها در اینجا آمده است:

• layout ( #۶۲۹ ).
• mhlo.frontend_attributes ( #628 ).
• mhlo.sharding ( #619 ).
• output_operand_aliases ( #740 ).
• فراداده مکان ( #594 ).

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

امضای عملیاتی شامل انواع تمام مقادیر ورودی (لیست انواع در سمت چپ -> ) و انواع تمام مقادیر خروجی (لیست انواع در سمت راست -> ) است. به طور دقیق، انواع ورودی زائد هستند و انواع خروجی نیز تقریباً همیشه زائد هستند (زیرا برای اکثر عملیات‌های StableHLO، انواع خروجی را می‌توان از ورودی‌ها استنباط کرد). با این وجود، امضای عملیاتی عمداً بخشی از نحو StableHLO برای سازگاری با MLIR است.

در زیر مثالی از یک تابع عملیاتی (op) آورده شده است که مخفف آن select_and_scatter است. این تابع ۳ مقدار ورودی ( %operand ، %source و %init_value )، ۲ تابع ورودی و ۳ ویژگی ورودی ( window_dimensions ، window_strides و padding ) را مصرف می‌کند. توجه داشته باشید که امضای تابع عملیاتی فقط شامل انواع مقادیر ورودی آن است (اما انواع توابع و ویژگی‌های ورودی که به صورت درون‌خطی ارائه می‌شوند را شامل نمی‌شود).

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

ثابت‌ها

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

ثابت‌های StableHLO دارای یک مقدار تحت‌اللفظی و یک نوع هستند که با هم یک مقدار StableHLO را نشان می‌دهند. به‌طورکلی، نوع بخشی از سینتکس ثابت است، مگر زمانی که مبهم نباشد (مثلاً یک ثابت بولی به‌طور واضح نوع i1 را دارد، در حالی که یک ثابت عدد صحیح می‌تواند چندین نوع ممکن داشته باشد).

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

ثابت‌های بولی مقادیر بولی true و false را نشان می‌دهند. ثابت‌های بولی از نوع i1 هستند.

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

ثابت‌های عدد صحیح، مقادیر صحیح را از طریق رشته‌هایی که از نمادگذاری اعشاری یا هگزادسیمال استفاده می‌کنند، نشان می‌دهند. سایر مبناها، مانند دودویی یا هشت هشتی، پشتیبانی نمی‌شوند. ثابت‌های عدد صحیح محدودیت‌های زیر را دارند:

• (C1) is_wellformed(integer_literal, integer_type) .

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

ثابت‌های ممیز شناور، مقادیر ممیز شناور را از طریق رشته‌هایی که از نمادگذاری اعشاری یا علمی استفاده می‌کنند، نشان می‌دهند. علاوه بر این، می‌توان از نمادگذاری هگزادسیمال برای مشخص کردن مستقیم بیت‌های اصلی در قالب ممیز شناور نوع مربوطه استفاده کرد. ثابت‌های ممیز شناور محدودیت‌های زیر را دارند:

• (C1) اگر از نمادگذاری غیر هگزادسیمال استفاده شود، is_wellformed(float_literal, float_type) کار می‌رود.
• (C2) اگر از نمادگذاری هگزادسیمال استفاده شود، size(hexadecimal_digits) = num_bits(float_type) / 4 .

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

ثابت‌های مختلط مقادیر مختلط را با استفاده از لیست‌هایی از یک بخش حقیقی (اول می‌آید) و یک بخش موهومی (دوم می‌آید) نشان می‌دهند. برای مثال، (1.0, 0.0) : complex<f32> نشان دهنده 1.0 + 0.0i است و (0.0, 1.0) : complex<f32> نشان دهنده 0.0 + 1.0i است. ترتیب ذخیره این بخش‌ها در حافظه، توسط پیاده‌سازی تعریف می‌شود. ثابت‌های مختلط محدودیت‌های زیر را دارند:

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

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

ثابت‌های تانسور مقادیر تانسور را با استفاده از لیست‌های تو در تو که از طریق نمادگذاری NumPy مشخص می‌شوند، نشان می‌دهند. برای مثال، dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> یک مقدار تانسور را با نگاشت زیر از اندیس‌ها به عناصر نشان می‌دهد: {0, 0} => 1 , {0, 2} => 3 {0, 1} => 2 , {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) '>'

ثابت‌های تانسوری کوانتیزه، مقادیر تانسوری کوانتیزه را با استفاده از همان نمادگذاری ثابت‌های تانسوری نشان می‌دهند، با عناصری که به عنوان ثابت‌هایی از نوع ذخیره‌سازی خود مشخص شده‌اند. ثابت‌های تانسوری کوانتیزه محدودیت‌های زیر را دارند:

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

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

لیترال‌های رشته‌ای شامل بایت‌هایی هستند که با استفاده از کاراکترهای ASCII و توالی‌های escape مشخص می‌شوند. آن‌ها کدگذاری-آگنوستیک هستند، بنابراین تفسیر این بایت‌ها از نظر پیاده‌سازی تعریف شده است. لیترال‌های رشته‌ای از نوع string هستند.

عملیات

عضلات شکم

معناشناسی

عملیات abs را به صورت عنصر به عنصر روی تانسور operand انجام می‌دهد و یک تانسور result تولید می‌کند. بسته به نوع عنصر، موارد زیر را انجام می‌دهد:

• برای اعداد صحیح علامت‌دار: مدول عدد صحیح.
• برای اعداد اعشاری: abs از IEEE-754.
• برای اعداد مختلط: مدول مختلط.
• برای انواع کوانتیزه: dequantize_op_quantize(abs, operand, type(result)) .

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (C1) shape(result) = shape(operand) .
• (C2) baseline_element_type(result) به صورت زیر تعریف می‌شود:
  • اگر is_complex( 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]

مثال‌های بیشتر

اضافه کردن

معناشناسی

جمع دو تانسور 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) . (یا: (is_per_axis_quantized(نتیجه)))
  • (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]]

مثال‌های بیشتر

بعد از همه

معناشناسی

تضمین می‌کند که عملیاتی که inputs تولید می‌کنند، قبل از هرگونه عملیاتی که به result وابسته است، اجرا شوند. اجرای این عملیات هیچ کاری انجام نمی‌دهد، فقط برای ایجاد وابستگی‌های داده‌ای از result به inputs وجود دارد.

ورودی‌ها

خروجی‌ها

مثال‌ها

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

مثال‌های بیشتر

همه_جمع_آوری

معناشناسی

در هر گروه فرآیند در شبکه فرآیند StableHLO، مقادیر تانسورهای operands را از هر فرآیند در امتداد all_gather_dim به هم پیوند می‌دهد و تانسورهای results تولید می‌کند.

این عملیات، شبکه فرآیند StableHLO را به process_groups تقسیم می‌کند که به شرح زیر تعریف می‌شوند:

• اگر channel_id <= 0 and use_global_device_ids = false 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 = true flattened_ids(replica_groups) .

پس از آن، درون هر 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) به صورت زیر تعریف می‌شود:
  • اگر از cross_replica استفاده شود، num_replicas .
  • اگر از cross_replica_and_partition استفاده شود، num_replicas .
  • اگر از flattened_ids استفاده شود، num_processes فعال می‌شود.
• (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 تقسیم می‌کند که به شرح زیر تعریف می‌شوند:

• اگر channel_id <= 0 and use_global_device_ids = false 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 = true flattened_ids(replica_groups) .

پس از آن، درون هر 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) به صورت زیر تعریف می‌شود:
  • اگر از cross_replica استفاده شود، num_replicas .
  • اگر از cross_replica_and_partition استفاده شود، num_replicas .
  • اگر از flattened_ids استفاده شود، num_processes فعال می‌شود.
• (C3) 0 <= replica_groups < size(replica_groups) .
• (C4) اگر use_global_device_ids = true باشد، آنگاه channel_id > 0 خواهد بود.
• computation (C5) از نوع (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]

مثال‌های بیشتر

همه_به_همه

معناشناسی

در هر گروه فرآیند در شبکه فرآیند StableHLO، مقادیر تانسورهای operands را در امتداد split_dimension به بخش‌هایی تقسیم می‌کند، بخش‌های تقسیم‌شده را بین فرآیندها پراکنده می‌کند، بخش‌های پراکنده را در امتداد concat_dimension به هم متصل می‌کند و تانسورهای results تولید می‌کند. این عملیات، شبکه فرآیند StableHLO را به process_groups تقسیم می‌کند که به صورت زیر تعریف می‌شوند:

• اگر channel_id <= 0 باشد، cross_replica(replica_groups)
• اگر channel_id > 0 ، cross_partition(replica_groups)

پس از آن، درون هر 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) به صورت زیر تعریف می‌شود:
  • اگر از cross_replica استفاده شود، num_replicas .
  • اگر از cross_partition استفاده شود، num_partitions .
• (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]]

مثال‌های بیشتر

آتان2

معناشناسی

عملیات 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 تولید می‌کند. به طور رسمی‌تر، این عملیات را می‌توان به صورت تجزیه به عملیات StableHLO موجود با استفاده از سینتکس پایتون به شرح زیر بیان کرد:

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 یکسانی هستند.
• operand (C3)، 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]

استنتاج هنجار_دسته‌ای

معناشناسی

تانسور operand در تمام ابعاد به جز بعد feature_index نرمال‌سازی می‌کند و یک تانسور result تولید می‌کند. به طور رسمی‌تر، این عملیات را می‌توان به صورت تجزیه به عملیات‌های StableHLO موجود با استفاده از سینتکس پایتون به صورت زیر بیان کرد:

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) .
• operand (C2)، 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 موجود با استفاده از سینتکس پایتون به شرح زیر بیان کرد:

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 returns in-memory representation of a given value, and its behavior is implementation-defined because the exact representation of tensors is implementation-defined, and the exact representation of element types is implementation-defined as well.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (C1) Given E = is_quantized(operand) ? storage_type(operand) : element_type(operand) , E' = is_quantized(result) ? storage_type(result) : element_type(result) , and R = rank(operand) :
  • If num_bits(E') = num_bits(E) , shape(result) = shape(operand) .
  • If num_bits(E') < num_bits(E) :
  • rank(result) = R + 1 .
  • dim(result, i) = dim(operand, i) for all 0 <= i < R .
  • dim(result, R) * num_bits(E') = num_bits(E) .
  • If num_bits(E') > num_bits(E) :
  • rank(result) = R - 1 .
  • dim(result, i) = dim(operand, i) for all 0 <= i < R .
  • dim(operand, R - 1) * num_bits(E) = num_bits(E') .
• (C2) If is_complex(operand) or is_complex(result) , then 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

معناشناسی

Expands the dimensions and/or rank of an input tensor by duplicating the data in the operand tensor and produces a result tensor. More formally, result[result_index] = operand[operand_index] where for all d in axes(operand) :

• operand_index[d] = 0 if dim(operand, d) = 1 .
• operand_index[d] = result_index[broadcast_dimensions[d]] otherwise.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) , scales(operand) , and zero_points(operand) may differ from quantization_dimension(result) , scales(result) , and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) For all d in axes(operand) :
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (C6) If is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • If dim(operand, quantization_dimension(operand)) = 1 , then 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]
//            ]
//          ]

مثال‌های بیشتر

مورد

معناشناسی

Produces the output from executing exactly one function from branches depending on the value of index . More formally, result = selected_branch() where:

• selected_branch = branches[index] if 0 <= index < size(branches) .
• selected_branch = branches[-1] otherwise.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

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

معناشناسی

Performs element-wise cubic root operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: rootn(x, 3) from IEEE-754.
• For complex numbers: complex cubic root.
• For quantized types: 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]

مثال‌های بیشتر

سقف

معناشناسی

Performs element-wise ceil of operand tensor and produces a result tensor. Implements the roundToIntegralTowardPositive operation from the IEEE-754 specification. For quantized types, performs 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

معناشناسی

تجزیه چولسکی دسته‌ای از ماتریس‌ها را محاسبه می‌کند.

More formally, for all i in index_space(result) , result[i0, ..., iR-3, :, :] is a Cholesky decomposition of a[i0, ..., iR-3, :, :] , in the form of either of a lower-triangular (if lower is true ) or upper-triangular (if lower is false ) matrix. The output values in the opposite triangle, ie the strict upper triangle or strict lower triangle correspondingly, are implementation-defined.

If there exists i where the input matrix is not an Hermitian positive-definite matrix, then the behavior is undefined.

For quantized types, performs 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]
//          ]

گیره

معناشناسی

Clamps every element of the operand tensor between a minimum and maximum value and produces a result tensor. More formally, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element) , where min_element = rank(min) = 0 ? min[] : min[result_index] , max_element = rank(max) = 0 ? max[] : max[result_index] . For quantized types, performs dequantize_op_quantize(clamp, min, operand, max, type(result)) .

Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers for this operation ( #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

معناشناسی

Within each process group in the StableHLO process grid, send the value of the operand tensor from the source process to the target processes and produce a result tensor.

The operation splits the StableHLO process grid into process_groups which is defined as follows:

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

Afterwards, result@process is given by:

• operand@process_groups[i, 0] if there exists an i such that the process is in process_groups[i] .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) otherwise.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (C1) is_unique(replica_groups) .
• (C2) 0 <= replica_groups < N where N is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (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

معناشناسی

Within each process group in the StableHLO process grid, sends the value of the operand tensor from the source process to the target process and produces a result tensor.

The operation splits the StableHLO process grid into process_groups which is defined as follows:

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

Afterwards, result@process is given by:

• operand@process_groups[i, 0] , if there exists an i such that process_groups[i, 1] = process .
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) otherwise.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (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 , where N is defined as:
  • num_replicas if cross_replica is used.
  • num_partitions if cross_partition is used.
• (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 را طبق comparison_direction و compare_type انجام می‌دهد و یک تانسور result تولید می‌کند.

The values of comparison_direction and compare_type have the following semantics:

For boolean and integer element types:

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

For floating-point element types with compare_type = FLOAT , the op implements the following IEEE-754 operations:

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

For floating-point element types with compare_type = TOTALORDER , the op uses the combination of totalOrder and compareQuietEqual operations from IEEE-754.

For complex element types, lexicographic comparison of (real, imag) pairs is performed using the provided comparison_direction and compare_type . Imposing an ordering on complex numbers involves surprising semantics, so in the future we are planning to remove support for complex numbers when comparison_direction is GE , GT , LE or LT ( #560 ).

For quantized types. performs 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 is defined as:
  • SIGNED if is_signed_integer(element_type(lhs)) .
  • UNSIGNED if is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)) .
  • FLOAT or TOTALORDER if is_float(element_type(lhs)) .
  • FLOAT if is_complex(element_type(lhs)) .

مثال‌ها

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %result: [true, false]

مثال‌های بیشتر

پیچیده

معناشناسی

Performs element-wise conversion to a complex value from a pair of real and imaginary values, lhs and rhs , and produces a result tensor.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (C1) type(lhs) = type(rhs) .
• (C2) shape(result) = shape(lhs) .
• (C3) element_type(result) has type complex<E> where 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)]

مثال‌های بیشتر

کامپوزیت

معناشناسی

Encapsulates an operation made up (composed) of other StableHLO operations, taking inputs and composite_attributes and producing results . The semantics of the op are implemented by the decomposition attribute. The composite op can be replaced with its decomposition without changing program semantics. In cases where inlining the decomposition does not provide the same op semantics, prefer using custom_call .

The version field (defaults to 0 ) is used to denote when a composite's semantics change.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

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

مثال‌های بیشتر

به هم پیوند دادن

معناشناسی

Concatenates inputs along dimension dimension in the same order as the given arguments and produces a result tensor. More formally, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1] , where:

1. id = d0 + ... + dk-1 + kd .
2. d is equal to dimension , and d0 , ... are d th dimension sizes of inputs .

ورودی‌ها

خروجی‌ها

محدودیت‌ها

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

مثال‌ها

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

مثال‌های بیشتر

ثابت

معناشناسی

Produces an output tensor from a constant 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]]

مثال‌های بیشتر

تبدیل کردن

معناشناسی

Performs an element-wise conversion from one element type to another on operand tensor and produces a result tensor.

For boolean-to-any-supported-type conversions, the value false is converted to zero, and the value true is converted to one. For any-supported-type-to-boolean conversions, a zero value is converted to false , and non-zero values are converted to true . See below for how this work for complex types.

For conversions involving integer-to-integer , integer-to-floating-point or floating-point-to-floating-point , if the source value can be exactly represented in the destination type, the result value is that exact representation. Otherwise, the behavior is TBD ( #180 ).

For conversions involving floating-point-to-integer , the fractional part is truncated. If the truncated value cannot be represented in the destination type, the behavior is TBD ( #180 ).

Conversion involving complex-to-complex follow the same behavior of floating-point-to-floating-point conversions for converting real and imaginary parts.

For complex-to-any-other-type and any-other-type-to-complex conversions, the source imaginary value is ignored or the destination imaginary value is zeroed, respectively. The conversion of the real part follows the floating-point conversions.

In principle, this operation could express dequantization (conversion from quantized tensors to regular tensors), quantization (conversion from regular tensors to quantized tensors) and requantization (conversion between quantized tensors), but at the moment we have dedicated operations for that - uniform_dequantize for the first use case and uniform_quantize for the second and the third use cases. In the future, these two ops may be merged into 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)]

مثال‌های بیشتر

کانولوشن

معناشناسی

Computes dot products between windows of lhs and slices of rhs and produces result . The following diagram shows how elements in result are computed from lhs and rhs using a concrete example.

More formally, consider the following reframing of the inputs in terms of lhs in order to be able to express windows of 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) .

This reframing uses the following helper functions:

• 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] where j[d] = i[permutation[d]] .

If feature_group_count = 1 and batch_group_count = 1 , then for all output_spatial_index in index_space(dim(result, output_spatial_dimensions...)) , result[result_shape(:, output_spatial_index, :)] = dot_product where:

• 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]) . This feature appears to be unused, so in the future we are planning to remove it ( #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]) .

If 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) .

If 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) .

For quantized types, performs 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)) .

For hybrid quantized types, performs 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) Given 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) Given 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) Given 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) is defined as:
  • dim(lhs, input_batch_dimension) / batch_group_count if result_dim = output_batch_dimension .
  • dim(rhs, kernel_output_feature_dimension) if result_dim = output_feature_dimension .
  • num_windows otherwise, where:
  • 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 .
• If the operation uses non-quantized tensors:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result) .
• If the operation uses quantized tensors:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C29) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) = kernel_output_feature_dimension .
  • (C30) If is_per_axis_quantized(result) , then quantization_dimension(result) = output_feature_dimension .
  • If is_quantized(lhs) :
  • (C31) storage_type(lhs) = storage_type(rhs) .
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C33) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !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]]
//          ]]

مثال‌های بیشتر

کسینوس

معناشناسی

Performs element-wise cosine operation on operand tensor and produces a result tensor. Depending on the element type, does the following:

• For floats: cos from IEEE-754.
• For complex numbers: complex cosine.
• For quantized types: 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

معناشناسی

Performs element-wise count of the number of leading zero bits in the operand tensor and produces a result tensor.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

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

معناشناسی

Encapsulates an implementation-defined operation call_target_name that takes inputs and called_computations and produces results . has_side_effect , backend_config and api_version may be used to provide additional implementation-defined metadata.

At the moment, this operation contains a fairly disorganized collection of metadata which reflects organic evolution of its counterpart operation in the XLA compiler. In the future, we are planning to unify this metadata ( #741 ).

ورودی‌ها

خروجی‌ها

(XLA GPU Support) Special custom_call targets

There are three special call_target_name related to buffer types: CreateBuffer creates an uninitialized buffer , Pin creates an initialized buffer and Unpin deallocates a buffer and returns the content of the buffer .

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

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin",
  api_version = 4 : i32,
} : (tensor<4xf64>) -> memref<4xf64>

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_name = "Unpin",
  api_version = 4 : i32,
} : (memref<4xf64>) -> tensor<4xf64>

نام مستعار

Some custom_call ops may require a part in the outputs and a part in the operands to share the same memory. This can be expressed via output_operand_aliases . An alias pair representation consists a list of output tuple indices representing the output part, and an operand_index along with a list of operand tuple indices representing the operand part. The list of output or operand tuple indices is empty if the corresponding type is not a tuple type, and can be arbitrarily long for an arbitrarily nested tuple type. This is similar to the XLA alias representation .

The output part and the input part in an alias pair must have the same type. For custom_call ops that aren't call to CreateBuffer , Pin and Unpin , a buffer operand can appear in at most one pair of alias, and a buffer output must appear in one pair of alias.

مثال‌ها

%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>

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases = [
    #stablehlo.output_operand_alias<output_tuple_indices = [],
      operand_index = 0,
      operand_tuple_indices = []>]
} : (memref<4xf64>) -> memref<4xf64>

تقسیم کردن

معناشناسی

Performs element-wise division of dividend lhs and divisor rhs tensors and produces a result tensor. Depending on the element type, does the following:

• For integers: integer division which produces the algebraic quotient with any fractional part discarded.
• For floats: division from IEEE-754.
• For complex numbers: complex division.
• For quantized types:
  • 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

معناشناسی

Computes dot products between slices of lhs and slices of rhs and produces a result tensor.

More formally, result[result_index] = dot_product , where:

• 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 where size(result_batching_index) = size(lhs_batching_dimensions) , size(result_lhs_index) = size(lhs_result_dimensions) and 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)) .

For quantized types, performs 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)) .

For hybrid quantized types, performs 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 controls the tradeoff between speed and accuracy for computations on accelerator backends. This can be one of the following (at the moment, the semantics of these enum values is underspecified, but we are planning to address this in #755 ):

• DEFAULT : Fastest calculation, but least accurate approximation to the original number.
• HIGH : Slower calculation, but more accurate approximation to the original number.
• HIGHEST : Slowest calculation, but most accurate approximation to the original number.

A DotAlgorithm defines the main properties of the algorithm used to implement the dot operation, which also defines the precision. If the algorithm attribute fields are set, then the precision_config must be DEFAULT . DotAlgorithms do not have a default value, as the default parameters are implementation defined. As such, all dot algorithm fields may be set to None to specify an empty dot algorithm, which will instead use the precision_config value.

DotAlgorithm fields include:

• lhs_precision_type and rhs_precision_type , the precisions that the LHS and RHS of the operation are rounded to. Precision types are independent from the storage types of the inputs and the output.
• accumulation_type the precision used for accumulation.
• lhs_component_count , rhs_component_count , and num_primitive_operations apply when we are doing an algorithm which decomposes the LHS and/or RHS into multiple components and does multiple "primitive" dot operations on those values - usually to emulate a higher precision (eg Leveraging the bfloat16 Artificial Intelligence Datatype For Higher-Precision Computations : bf16_6x tf32_3x, etc). For algorithms with no decomposition, these values should be set to 1 .
• allow_imprecise_accumulation to specify if accumulation in lower precision is permitted for some steps (eg CUBLASLT_MATMUL_DESC_FAST_ACCUM ).

Example DotAlgorithm attributes:

// 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}

It is up to the implementations to decide which combinations are supported. In general, it is not guaranteed that each algorithm is supported on each accelerator type by the consumer of the StableHLO. If a given algorithm is not supported, an error should be raised as opposed to falling back to an alternative. StableHLO verification will provide best effort verification, preventing algorithms that are not known to be supported on any hardware.

See xla_data.proto > Algorithm for some supported algorithm values. Ticket #2483 captures the plan to create a centralized doc on supported algorithms by backend.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (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) .
• If the operation uses non-quantized tensors:
  • (C13) element_type(lhs) = element_type(rhs) .
• If the operation uses quantized tensors:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C15) zero_points(rhs) = 0 .
  • (C16) If is_per_axis_quantized(rhs) , then quantization_dimension(rhs) not in rhs_contracting_dimensions .
  • If is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs) .
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C19) If is_per_tensor_quantized(rhs) , then is_per_tensor_quantized(result) .
  • If !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result) .
• If !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

معناشناسی

This operation is functionally identical to broadcast_in_dim op, but the result shape is specified dynamically via output_dimensions .

The operation also accepts optional attributes known_expanding_dimensions , known_nonexpanding_dimensions to express static knowledge about the expanding behavior of dimensions. If not specified, all dimensions are assumed to be possibly expanding.

ورودی‌ها

خروجی‌ها

محدودیت‌ها

• (C1) element_type(result) is given by:
  • element_type(operand) , if !is_per_axis_quantized(operand) .
  • element_type(operand) except that quantization_dimension(operand) , scales(operand) , and zero_points(operand) may differ from quantization_dimension(result) , scales(result) , and zero_points(result) resp., otherwise.
• (C2) size(broadcast_dimensions) = rank(operand) .
• (C3) 0 <= broadcast_dimensions < rank(result) .
• (C4) is_unique(broadcast_dimensions) .
• (C5) For all d in axes(operand) :
  • dim(operand, d) = 1 or
  • dim(operand, d) = dim(result, broadcast_dimensions[d]) .
• (C6) If is_per_axis_quantized(result) :
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)] .
  • If dim(operand, quantization_dimension(operand)) = 1 , then 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

معناشناسی

This operation is functionally identical to convolution op, but the padding is specified dynamically via padding .

ورودی‌ها