مشخصات StableHLO

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

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

این مشخصات شامل سه بخش عمده است. ابتدا، بخش برنامه‌ها ساختار برنامه‌های StableHLO را توصیف می‌کند که از توابع StableHLO تشکیل شده‌اند که خود از عملیات StableHLO تشکیل شده‌اند. در آن ساختار، بخش Ops معنای عملیات های فردی را مشخص می کند. بخش Execution معنایی را برای همه این عملیات ها که با هم در یک برنامه اجرا می شوند ارائه می دهد. در نهایت، بخش Notation در مورد نماد استفاده شده در تمام مشخصات بحث می کند.

برای مشاهده مشخصات نسخه قبلی StableHLO، مخزن را در نسخه مورد علاقه با برچسب باز کنید. به عنوان مثال، StableHLO v0.19.0 Spec . برای مشاهده تغییراتی که در هر برآمدگی جزئی نسخه StableHLO رخ داده است، به گزارش نسخه در VhloDialect.td مراجعه کنید.

برنامه ها

Program ::= {Func}

برنامه های StableHLO از تعداد دلخواه توابع StableHLO تشکیل شده اند. در زیر یک برنامه نمونه با تابع @main وجود دارد که دارای 3 ورودی ( %image ، %weights و %bias ) و 1 خروجی است. بدنه تابع دارای 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> یک نوع تانسور با شکل 2x3 و نوع عنصر f32 است. دارای دو بعد (یا به عبارتی دو محور) - بعد 0 و بعد 1 - که اندازه آنها 2 و 3 است. رتبه آن 2 است.

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

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

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

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

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

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

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

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

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

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

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

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

TokenType ::= 'token'

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

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

انواع تاپی نشان دهنده تاپل ها، یعنی لیست های ناهمگن هستند. Tuples یک ویژگی قدیمی است که فقط برای سازگاری با 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 با مقادیر تانسور 0 بعدی از نوع tensor<T> اصطلاحی است).

• نوع Boolean مقادیر بولی true و false را نشان می دهد.
• انواع عدد صحیح می توانند علامت دار ( si ) یا بدون علامت ( ui ) باشند و یکی از عرض بیت های پشتیبانی شده ( 2 ، 4 ، 8 ، 16 ، 32 یا 64 ) را داشته باشند. انواع siN علامت دار مقادیر صحیح را از -2^(N-1) تا 2^(N-1)-1 شامل، و انواع unsigned uiN مقادیر صحیح از 0 تا 2^N-1 را نشان می دهند.
• انواع ممیز شناور می تواند یکی از موارد زیر باشد:
  • اعداد ممیز شناور 8 بیتی f8E3M4 ، f8E4M3 و f8E5M2 طبق قراردادهای IEEE-754.
  • انواع f8E4M3FN و f8E5M2 مربوط به کدگذاری های E4M3 و E5M2 فرمت FP8 که در فرمت های FP8 برای یادگیری عمیق توضیح داده شده است.
  • انواع f8E4M3FNUZ و f8E5M2FNUZ مربوط به کدگذاری های E4M3 و E5M2 فرمت های FP8 که در قالب های عددی 8 بیتی برای شبکه های عصبی عمیق توضیح داده شده است.
  • نوع f8E4M3B11FNUZ مربوط به کدگذاری E4M3 فرمت‌های FP8 که در آموزش و استنتاج هیبریدی 8 بیتی شناور نقطه شناور (HFP8) برای شبکه‌های عصبی عمیق توضیح داده شده است.
  • نوع bf16 مربوط به فرمت bfloat16 شرح داده شده در BFloat16: راز عملکرد بالا در Cloud TPU .
  • انواع f16 ، f32 و f64 مربوط به فرمت‌های باینری binary16 ("نیمه دقت")، binary32 ("دقت تک") و binary64 ("دقت دوگانه") که در استاندارد IEEE 754 توضیح داده شده‌اند.
  • نوع tf32 با فرمت TensorFloat32 مطابقت دارد و در StableHLO پشتیبانی محدودی دارد.
  • انواع f4E2M1FN ، f6E2M3FN ، f6E3M2FN و f8E8M0FNU MX (ریزمقیاس‌سازی) در مشخصات فرمت‌های OCP Microscaling توضیح داده شده‌اند.
• انواع مختلط مقادیر پیچیده ای را نشان می دهند که یک قسمت واقعی و یک قسمت خیالی از همان نوع عنصر دارند. انواع پیچیده پشتیبانی شده complex<f32> (هر دو قسمت از نوع f32 هستند) و complex<f64> (هر دو قسمت از نوع f64 هستند).

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

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

StringType ::= 'string'

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

عملیات

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

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

عملیات StableHLO (که به آنها ops نیز می گویند) دارای یک نام، ورودی/خروجی و یک امضا هستند. نام از stablehlo. پیشوند و یادگاری که به طور منحصر به فرد یکی از عملیات های پشتیبانی شده را شناسایی می کند. برای یک لیست جامع از همه عملیات های پشتیبانی شده به زیر مراجعه کنید.

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

عملیات ها ورودی ها را مصرف می کنند و خروجی ها را تولید می کنند. ورودی ها به مقادیر ورودی (محاسبه شده در طول اجرا)، توابع ورودی (به صورت ایستا ارائه می شوند، زیرا در StableHLO توابع مقادیر درجه یک نیستند) و ویژگی های ورودی (همچنین به صورت ایستا ارائه می شوند) دسته بندی می شوند. نوع ورودی و خروجی های مصرف شده و تولید شده توسط یک op به حافظه آن بستگی دارد. برای مثال، add op 2 مقدار ورودی مصرف می کند و 1 مقدار خروجی تولید می کند. در مقایسه، گزینه select_and_scatter 3 مقدار ورودی، 2 تابع ورودی و 3 ویژگی ورودی مصرف می کند.

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

توابع ورودی (که توابع ناشناس نیز نامیده می شوند) بسیار شبیه توابع نامگذاری شده هستند با این تفاوت که: 1) آنها شناسه ندارند (از این رو نام "ناشناس")، 2) انواع خروجی را اعلام نمی کنند (انواع خروجی عبارتند از استنتاج از عملیات return در تابع).

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

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

ویژگی های ورودی یک نام و یک مقدار دارند که یکی از ثابت های پشتیبانی شده است. آنها راه اصلی برای تعیین ابرداده استاتیک برای عناصر برنامه هستند. به عنوان مثال، concatenate op dimension ویژگی برای مشخص کردن بعدی که مقادیر ورودی آن در امتداد آن الحاق می شود، استفاده می کند. به طور مشابه، slice op از چندین ویژگی مانند 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}] ')'

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

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

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

ثابت ها

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

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

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

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

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

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

• (C1) is_wellformed(integer_literal, integer_type) .

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

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

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

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

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

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

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

ثابت های تانسور مقادیر تانسور را با استفاده از لیست های تودرتو مشخص شده از طریق نماد NumPy نشان می دهند. برای مثال، dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> یک مقدار تانسور را با نگاشت زیر از شاخص ها به عناصر نشان می دهد: {0, 0} => 1 , {0, 1} => 2 ، {0, 2} => 3 ، {1, 0} => 4 ، {1, 1} => 5 , {1, 2} => 6 . ترتیبی که این عناصر سپس در حافظه ذخیره می شوند به صورت پیاده سازی تعریف شده است. ثابت های تانسور دارای محدودیت های زیر هستند:

• (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 و دنباله های فرار مشخص شده اند. آنها کدگذاری-آگنوستیک هستند، بنابراین تفسیر این بایت ها به صورت پیاده سازی تعریف شده است. حرفهای رشته ای دارای نوع string هستند.

عملیات

عضلات شکم

معناشناسی

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

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

ورودی ها

خروجی ها

محدودیت ها

• (C1) shape(result) = shape(operand) .
• (C2) baseline_element_type(result) به صورت زیر تعریف می شود:
  • complex_element_type(element_type(operand)) if is_complex(operand) است.
  • baseline_element_type(operand) در غیر این صورت.

نمونه ها

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

نمونه های بیشتر

اضافه کردن

معناشناسی

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

• برای Booleans: منطقی OR.
• برای اعداد صحیح: جمع اعداد صحیح.
• برای شناورها: addition از IEEE-754.
• برای اعداد مختلط: جمع مختلط.
• برای انواع کوانتیزه شده: dequantize_op_quantize(add, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

• اگر عملیات از تانسورهای غیر کوانتیزه استفاده می کند:
  • (C1) type(lhs) = type(rhs) = type(result) .
• اگر عملیات از تانسورهای کوانتیزه استفاده می کند:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result) .
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result) .
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result) .
  • (C6) اگر is_per_axis_quantized(lhs) ، سپس quantization_dimension(lhs) = quantization_dimension(result) .
  • (C7) اگر is_per_axis_quantized(rhs) ، سپس quantization_dimension(rhs) = quantization_dimension(result) .

نمونه ها

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

نمونه های بیشتر

after_all

معناشناسی

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

ورودی ها

خروجی ها

نمونه ها

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

نمونه های بیشتر

همه_جمع آوری

معناشناسی

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

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

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

پس از آن، در هر process_group :

• operands...@receiver = [operand@sender for sender in process_group] برای همه receiver در process_group .
• results...@process = concatenate(operands...@process, all_gather_dim) برای همه process در process_group .

ورودی ها

خروجی ها

محدودیت ها

• (C1) 0 <= all_gather_dim < rank(operands...) .
• (C2) is_unique(replica_groups) .
• size(replica_groups) به صورت زیر تعریف می شود:
  • num_replicas اگر cross_replica استفاده شود.
  • num_replicas اگر cross_replica_and_partition استفاده شود.
  • اگر از 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 تقسیم می‌کند که به صورت زیر تعریف می‌شوند:

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

پس از آن، در هر process_group :

• results...@process[result_index] = exec(schedule) برای برخی از schedule درختی باینری که در آن:
  • exec(node) = computation(exec(node.left), exec(node.right)) .
  • exec(leaf) = leaf.value .
• schedule یک درخت باینری تعریف شده توسط پیاده سازی است که پیمایش ترتیبی آن to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])) .

ورودی ها

خروجی ها

محدودیت ها

• (C1) is_unique(replica_groups) .
• size(replica_groups) به صورت زیر تعریف می شود:
  • num_replicas اگر cross_replica استفاده شود.
  • num_replicas اگر cross_replica_and_partition استفاده شود.
  • اگر از flattened_ids استفاده شود، num_processes .
• (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) if 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) .
• 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 تولید می کند. بسته به نوع عنصر، موارد زیر را انجام می دهد:

• برای Booleans: منطقی و.
• برای اعداد صحیح: بیتی 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 موجود با استفاده از نحو 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]

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

معناشناسی

تانسور 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]]
//          ]

آموزش هنجارهای دسته ای

معناشناسی

میانگین و واریانس را در تمام ابعاد به جز بعد 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

معناشناسی

یک عملیات بیت‌کست را روی تانسور 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]
//            ]
//          ]

نمونه های بیشتر

مورد

معناشناسی

بسته به مقدار index خروجی را از اجرای دقیقاً یک عملکرد از branches تولید می کند. به طور رسمی ، 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 ای تولید می کند. بسته به نوع عنصر ، موارد زیر را انجام می دهد:

• برای شناورها: 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]

نمونه های بیشتر

سقف

معناشناسی

مجاری عناصر عناصر تانسور 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]

نمونه های بیشتر

بی پروا

معناشناسی

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

به طور رسمی تر ، برای همه i در index_space(result) ، result[i0, ..., iR-3, :, :] یک تجزیه چولسکی از a[i0, ..., iR-3, :, :] به شکل هر یک از ماتریس های مثلثی پایین تر (اگر lower true ) یا مثلثی فوقانی (اگر lower false باشد). مقادیر خروجی در مثلث مخالف ، یعنی مثلث فوقانی بالا یا مثلث پایین سخت به طور متناوب ، اجرا شده است.

اگر وجود داشته باشد که i در جایی که ماتریس ورودی یک ماتریس دینهای مثبت هرمیتی نیست ، رفتار تعریف نشده است.

برای انواع کمیت ، dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)) را انجام می دهد.

ورودی ها

خروجی ها

محدودیت ها

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

نمونه ها

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

گیره

معناشناسی

هر عنصر از تانسور operand بین حداقل و حداکثر مقدار گیره می کند و تانسور result را تولید می کند. به طور رسمی تر ، result[result_index] = minimum(maximum(operand[result_index], min_element), max_element) ، جایی که min_element = rank(min) = 0 ? min[] : min[result_index] ، max_element = rank(max) = 0 ? max[] : max[result_index] . برای انواع کم ، dequantize_op_quantize(clamp, min, operand, max, type(result)) را انجام می دهد.

تحمیل سفارش برای اعداد پیچیده شامل معناشناسی شگفت آور است ، بنابراین در آینده قصد داریم پشتیبانی از شماره های پیچیده را برای این عملیات حذف کنیم ( شماره 560 ).

ورودی ها

خروجی ها

محدودیت ها

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

نمونه ها

// %min: [5, 10, 15]
// %operand: [3, 13, 23]
// %max: [10, 15, 20]
%result = "stablehlo.clamp"(%min, %operand, %max) : (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> tensor<3xi32>
// %result: [5, 13, 20]

نمونه های بیشتر

مجموعه_

معناشناسی

در هر گروه فرآیند در شبکه فرآیند 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]]

مجموعه_پرموت

معناشناسی

در هر گروه فرآیند در شبکه فرآیند 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]]

نمونه های بیشتر

مقایسه کنید

معناشناسی

مقایسه عناوین عناصر تنشور 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 ، OP عملیات IEEE-754 زیر را پیاده سازی می کند:

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

برای انواع عنصر نقطه شناور با compare_type = TOTALORDER ، OP از ترکیبی از 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 به صورت تعریف شده است:
  • اگر is_signed_integer(element_type(lhs)) SIGNED .
  • 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 تولید می کند. معناشناسی OP توسط ویژگی decomposition اجرا می شود. composite OP را می توان با تجزیه آن بدون تغییر معانی برنامه جایگزین کرد. در مواردی که در حال تجزیه تجزیه ، همان معانی OP را ارائه نمی دهد ، استفاده از 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 را انجام می دهد و یک تانسور result تولید می کند.

برای تبدیل های بولی به هر نوع پشتیبانی ، مقدار false به صفر تبدیل می شود و مقدار true به یک تبدیل می شود. برای تبدیل های هر نوع پشتیبانی از نوع به بولایی ، یک مقدار صفر به مقادیر false تبدیل می شود و مقادیر غیر صفر به true تبدیل می شوند. در زیر نحوه عملکرد این کار برای انواع پیچیده را مشاهده کنید.

برای تبدیل های مربوط به عدد صحیح ، عدد صحیح به شناور یا نقطه شناور نقطه به شناور ، اگر مقدار منبع دقیقاً در نوع مقصد نشان داده شود ، مقدار نتیجه آن نمایش دقیق است. در غیر این صورت ، رفتار TBD است ( شماره 180 ).

برای تبدیل های مربوط به شناور-نقطه به دین ، ​​قسمت کسری کوتاه می شود. اگر مقدار کوتاه شده در نوع مقصد قابل نمایش نباشد ، رفتار TBD است ( #180 ).

تبدیل شامل پیچیده به پیچیده ، از همان رفتار تبدیل های نقطه شناور به شناور برای تبدیل قطعات واقعی و خیالی پیروی می کند.

برای تبدیل های پیچیده به هر نوع و هر نوع دیگری از نوع دیگر ، ارزش تخیلی منبع نادیده گرفته می شود یا مقدار تخیل مقصد به ترتیب صفر می شود. تبدیل قسمت واقعی از تبدیل نقطه شناور پیروی می کند.

در اصل ، این عملیات می تواند بیان مجدد (تبدیل از تنسورهای کمیت به تانسرهای منظم) ، کمیت (تبدیل از تنسورهای منظم به تنسورهای کمیت) و نیاز (تبدیل بین تنسورهای کمیت) ، اما در حال حاضر ما عملیات اختصاصی برای آن داریم - یکنواخت برای آن - 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)]

نمونه های بیشتر

پیچیدگی

معناشناسی

محصولات DOT را بین ویندوز 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)) .

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) با توجه به input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension] :
  • is_unique(input_dimensions) .
  • 0 <= input_dimensions < N .
• 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 ای را تولید می کند. بسته به نوع عنصر ، موارد زیر را انجام می دهد:

• برای شناورها: 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]]

نمونه های بیشتر

CUSTION_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 و تانسور Divisor rhs را انجام می دهد و تانسور result ای را تولید می کند. بسته به نوع عنصر ، موارد زیر را انجام می دهد:

• برای عدد صحیح: تقسیم عدد صحیح که میزان جبری را با هر قسمت کسری دور انداخته می شود.
• برای شناورها: division از IEEE-754.
• برای اعداد پیچیده: تقسیم پیچیده.
• برای انواع کمیت:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)) .

ورودی ها

خروجی ها

محدودیت ها

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

نمونه ها

// %lhs: [17.1, -17.1, 17.1, -17.1]
// %rhs: [3.0, 3.0, -3.0, -3.0]
%result = "stablehlo.divide"(%lhs, %rhs) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// %result: [5.66666651, -5.66666651, -5.66666651, 5.66666651]

نمونه های بیشتر

dot_general

معناشناسی

محصولات DOT را بین برش های 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)) انجام می دهد. 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 تجارت بین سرعت و دقت محاسبات در باکتری های شتاب دهنده را کنترل می کند. این می تواند یکی از موارد زیر باشد (در حال حاضر ، معناشناسی این مقادیر enum مشخص نشده است ، اما ما در حال برنامه ریزی برای پرداختن به این موضوع در شماره 755 هستیم):

• DEFAULT : سریعترین محاسبه ، اما حداقل تقریب دقیق به شماره اصلی.
• HIGH : محاسبه آهسته تر ، اما تقریب دقیق تر به شماره اصلی.
• HIGHEST : کمترین محاسبه ، اما دقیق ترین تقریب به تعداد اصلی.

یک DotAlgorithm خصوصیات اصلی الگوریتم مورد استفاده برای اجرای عملیات DOT را تعریف می کند ، که همچنین دقت را تعریف می کند. اگر فیلدهای ویژگی الگوریتم تنظیم شده باشد ، باید precision_config DEFAULT باشد. DotAlgorithms مقدار پیش فرض ندارند ، زیرا پارامترهای پیش فرض پیاده سازی تعریف شده است. به این ترتیب ، تمام زمینه های الگوریتم DOT ممکن است روی None تنظیم شوند تا یک الگوریتم نقطه خالی را مشخص کنند ، که در عوض از مقدار precision_config استفاده می کند.

زمینه های DotAlgorithm شامل موارد زیر است:

• lhs_precision_type و rhs_precision_type ، مواردی که LHS و RHS این عملیات گرد شده اند. انواع دقیق مستقل از انواع ذخیره ورودی ها و خروجی هستند.
• accumulation_type دقت مورد استفاده برای تجمع.
• lhs_component_count ، rhs_component_count ، و num_primitive_operations هنگامی که ما یک الگوریتم را انجام می دهیم که LHS و/یا RHS را در چندین مؤلفه تجزیه می کند ، اعمال می شود و چندین عمل "بدوی" را بر روی آن مقادیر انجام می دهد - معمولاً برای تقلید از یک نمونه گیری بالاتر (به عنوان مثال ، با استفاده از Bfloat16 Artifical Artifical Artificial Artifical برای دقت بالاتر محاسبات : 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 برنامه ایجاد یک سند متمرکز را در الگوریتم های پشتیبانی شده توسط 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) .
• اگر این عملیات از تانسور غیرقانونی استفاده می کند:
  • (C13) element_type(lhs) = element_type(rhs) .
• در صورت استفاده از تانسور کمیته:
  • (c14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs) .
  • (C15) zero_points(rhs) = 0 .
  • (C16) اگر is_per_axis_quantized(rhs) ، سپس quantization_dimension(rhs) در rhs_contracting_dimensions نیست.
  • اگر is_quantized(lhs) :
  • (C17) storage_type(lhs) = storage_type(rhs) .
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result) .
  • (C19) اگر is_per_tensor_quantized(rhs) ، سپس is_per_tensor_quantized(result) .
  • اگر !is_quantized(lhs) :
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result) .
• اگر !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation) :
  • (C21) precision_config... = DEFAULT .
  • (C22) 0 < lhs_component_count .
  • (C23) 0 < rhs_component_count .
  • (C24) 0 < num_primitive_operations .

نمونه ها

// %lhs: [
//        [[1, 2],
//         [3, 4]],
//        [[5, 6],
//         [7, 8]]
//       ]
// %rhs: [
//        [[1, 0],
//         [0, 1]],
//        [[1, 0],
//         [0, 1]]
//       ]
%result = "stablehlo.dot_general"(%lhs, %rhs) {
  dot_dimension_numbers = #stablehlo.dot<
    lhs_batching_dimensions = [0],
    rhs_batching_dimensions = [0],
    lhs_contracting_dimensions = [2],
    rhs_contracting_dimensions = [1]
  >,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
  algorithm = #stablehlo.dot_algorithm<
    lhs_precision_type = tf32,
    rhs_precision_type = tf32,
    accumulation_type = f32,
    lhs_component_count = 1,
    rhs_component_count = 1,
    num_primitive_operations = 1,
    allow_imprecise_accumulation = false
  >
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

نمونه های بیشتر

dynamic_broadcast_in_dim

معناشناسی

این عمل از نظر عملکردی با Broadcast_in_dim OP یکسان است ، اما شکل نتیجه از طریق output_dimensions به صورت پویا مشخص می شود.

این عملیات همچنین ویژگی های اختیاری را می known_expanding_dimensions ، known_nonexpanding_dimensions برای بیان دانش استاتیک در مورد رفتار در حال گسترش ابعاد. اگر مشخص نشده باشد ، فرض بر این است که همه ابعاد در حال گسترش هستند.

ورودی ها

خروجی ها

محدودیت ها

• (C1) element_type(result) توسط:
  • 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 .

ورودی ها

خروجی ها

محدودیت ها

• (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]]]
//        ]
// %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]]
//          ]]