Spesifikasi StableHLO

StableHLO adalah set operasi untuk operasi tingkat tinggi (HLO) dalam model machine learning (ML). StableHLO berfungsi sebagai lapisan portabilitas antara berbagai framework ML dan compiler ML: framework ML yang menghasilkan program StableHLO kompatibel dengan compiler ML yang menggunakan program StableHLO.

Tujuan kami adalah menyederhanakan dan mempercepat pengembangan ML dengan menciptakan lebih banyak interoperabilitas antara berbagai framework ML (seperti TensorFlow, JAX, dan PyTorch) dan compiler ML (seperti XLA dan IREE). Untuk itu, dokumen ini memberikan spesifikasi untuk bahasa pemrograman StableHLO.

Spesifikasi ini berisi tiga bagian utama. Pertama, bagian Program menjelaskan struktur program StableHLO yang terdiri dari fungsi StableHLO yang terdiri dari operasi StableHLO. Dalam struktur tersebut, bagian Ops menentukan semantik setiap operasi. Bagian Execution menyediakan semantik untuk semua operasi ini yang dijalankan bersama dalam sebuah program. Terakhir, bagian Notasi membahas notasi yang digunakan di seluruh spesifikasi.

Untuk melihat spesifikasi dari rilis StableHLO sebelumnya, buka repo di rilis yang diberi tag yang diinginkan. Misalnya, Spesifikasi StableHLO v0.19.0. Untuk melihat perubahan yang terjadi pada setiap peningkatan versi minor StableHLO, lihat log versi di VhloDialect.td.

Program

Program ::= {Func}

Program StableHLO terdiri dari sejumlah fungsi StableHLO. Di bawah ini adalah contoh program dengan fungsi @main yang memiliki 3 input (%image, %weights, dan %bias) serta 1 output. Isi fungsi memiliki 6 operasi.

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>) -> ()
}

Fungsi

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

Fungsi StableHLO (yang juga disebut fungsi bernama) memiliki ID, input/output, dan isi. Pada masa mendatang, kami berencana untuk memperkenalkan metadata tambahan untuk fungsi guna mencapai kompatibilitas yang lebih baik dengan HLO (#425, #626, #740, #744).

ID

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

ID StableHLO mirip dengan ID di banyak bahasa pemrograman, dengan dua keunikan: 1) semua ID memiliki sigil yang membedakan berbagai jenis ID, 2) ID nilai dapat berupa angka sepenuhnya untuk menyederhanakan pembuatan program StableHLO.

Jenis

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

Jenis StableHLO dikategorikan ke dalam jenis nilai (yang juga disebut jenis kelas satu) yang merepresentasikan nilai StableHLO dan jenis non-nilai yang menjelaskan elemen program lainnya. Jenis StableHLO mirip dengan jenis di banyak bahasa pemrograman, dengan keunikan utama adalah sifat khusus domain StableHLO yang menghasilkan beberapa hasil yang tidak biasa (misalnya, jenis skalar bukan jenis nilai).

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

Jenis tensor merepresentasikan tensor, yaitu array multidimensi. Objek ini memiliki bentuk dan jenis elemen, dengan bentuk yang merepresentasikan ukuran dimensi non-negatif atau tidak diketahui dalam urutan menaik dari dimensi yang sesuai (yang juga disebut sumbu) yang diberi nomor dari 0 hingga R-1. Jumlah dimensi R disebut rank. Misalnya, tensor<2x3xf32> adalah jenis tensor dengan bentuk 2x3 dan jenis elemen f32. Array ini memiliki dua dimensi (atau, dengan kata lain, dua sumbu) - dimensi ke-0 dan dimensi ke-1 - yang ukurannya adalah 2 dan 3. Peringkatnya adalah 2.

Bentuk dapat sebagian atau sepenuhnya tidak diketahui (dinamis), misalnya, tensor<?x2xf64> sebagian tidak diketahui dan tensor<?x?xf64> sepenuhnya tidak diketahui. Ukuran dimensi dinamis ditampilkan menggunakan ?. Bentuk tidak dapat dibatalkan peringkatnya.

Pada masa mendatang, kami berencana untuk memperluas jenis tensor di luar ukuran dimensi dan jenis elemen, misalnya, untuk menyertakan tata letak (#629) dan kejarangan (#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

Jenis elemen terkuantisasi merepresentasikan nilai bilangan bulat dari jenis penyimpanan dalam rentang dari storage_min hingga storage_max (inklusif) yang sesuai dengan nilai floating point dari jenis yang dinyatakan. Untuk nilai bilangan bulat i tertentu, nilai floating point f yang sesuai dapat dihitung sebagai f = (i - zero_point) * scale, dengan scale dan zero_point disebut parameter kuantisasi. storage_min dan storage_max bersifat opsional dalam tata bahasa, tetapi memiliki nilai default min_value(storage_type) dan max_value(storage_type). Jenis elemen terkuantisasi memiliki batasan berikut:

• (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) Jika is_empty(quantization_dimension), maka size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Saat ini, QuantizationScale adalah konstanta floating point, tetapi ada minat yang kuat terhadap skala berbasis bilangan bulat, yang diwakili dengan pengali dan pergeseran. Kami berencana untuk menjelajahi fitur ini dalam waktu dekat (#1404).

Diskusi tentang semantik QuantizationZeroPoint sedang berlangsung, termasuk jenis, nilai, dan apakah hanya ada satu atau beberapa titik nol yang berpotensi ada dalam jenis tensor terkuantisasi. Berdasarkan hasil diskusi ini, spesifikasi seputar nol poin dapat berubah di masa mendatang (#1405).

Diskusi berkelanjutan lainnya melibatkan semantik QuantizationStorageMin dan QuantizationStorageMax untuk menentukan apakah ada batasan yang harus diterapkan pada nilai ini dan pada nilai tensor terkuantisasi (#1406).

Terakhir, kami berencana untuk mengeksplorasi cara menampilkan skala yang tidak diketahui dan titik nol, serupa dengan cara kami berencana untuk mengeksplorasi cara menampilkan ukuran dimensi yang tidak diketahui (#1407).

Jenis tensor terkuantisasi merepresentasikan tensor dengan elemen terkuantisasi. Tensor ini sama persis dengan tensor biasa, kecuali elemennya memiliki jenis elemen terkuantisasi, bukan jenis elemen biasa.

Dalam tensor terkuantisasi, kuantisasi dapat dilakukan per-tensor, yang berarti memiliki satu scale dan zero_point untuk seluruh tensor atau dapat dilakukan per-axis, yang berarti memiliki beberapa scales dan zero_points, satu pasangan per slice dari dimensi quantization_dimension tertentu. Secara lebih formal, dalam tensor t dengan kuantisasi per sumbu, ada dim(t, quantization_dimension) slice dari quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], dll. Semua elemen dalam slice ke-i menggunakan scales[i] dan zero_points[i] sebagai parameter kuantisasinya. Jenis tensor terkuantisasi memiliki batasan berikut:

• Untuk kuantisasi per-tensor:
  • Tidak ada batasan tambahan.
• Untuk kuantisasi per sumbu:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Jenis token merepresentasikan token, yaitu nilai buram yang dihasilkan dan digunakan oleh beberapa operasi. Token digunakan untuk memaksakan urutan eksekusi pada operasi seperti yang dijelaskan di bagian Execution.

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

Jenis buffer merepresentasikan buffer. Misalnya, di XLA, buffer adalah array multidimensi dengan penyimpanan yang konsisten. Mirip dengan jenis tensor, jenis buffer memiliki bentuk dan jenis elemen, dengan bentuk yang merepresentasikan ukuran dimensi non-negatif atau tidak diketahui dalam urutan menaik dari dimensi yang sesuai (yang juga disebut sumbu) yang diberi nomor dari 0 hingga R-1. Jumlah dimensi R disebut rank. Misalnya, memref<2x3xf32> adalah jenis buffer dengan bentuk 2x3 dan jenis elemen f32. Array ini memiliki dua dimensi (atau, dengan kata lain, dua sumbu) - dimensi ke-0 dan dimensi ke-1 - yang ukurannya adalah 2 dan 3. Peringkatnya adalah 2.

Buffer dapat dialokasikan menggunakan custom_call hingga CreateBuffer atau Pin dan dibatalkan alokasinya melalui custom_call hingga Unpin. Hanya operasi custom_call yang dapat membaca dan menulis konten di dalam buffer. Lihat custom_call untuk mengetahui detail selengkapnya.

Jenis tuple merepresentasikan tuple, yaitu daftar heterogen. Tuple adalah fitur lama yang hanya ada untuk kompatibilitas dengan HLO. Di HLO, tuple digunakan untuk merepresentasikan input dan output variadik. Di StableHLO, input dan output variadik didukung secara native, dan satu-satunya penggunaan tuple di StableHLO adalah untuk merepresentasikan ABI HLO secara komprehensif, misalnya T, tuple<T>, dan tuple<tuple<T>> mungkin sangat berbeda bergantung pada implementasi tertentu. Pada masa mendatang, kami berencana melakukan perubahan pada ABI HLO yang dapat memungkinkan kami menghapus jenis tuple dari 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'

Jenis elemen merepresentasikan elemen jenis tensor. Tidak seperti di banyak bahasa pemrograman, jenis ini bukan kelas utama di StableHLO. Artinya, program StableHLO tidak dapat secara langsung merepresentasikan nilai jenis ini (sebagai hasilnya, nilai skalar jenis T direpresentasikan secara idiomatis dengan nilai tensor 0 dimensi jenis tensor<T>).

• Jenis boolean merepresentasikan nilai boolean true dan false.
• Jenis bilangan bulat dapat ditandai (si) atau tidak ditandai (ui) dan memiliki salah satu lebar bit yang didukung (2, 4, 8, 16, 32, atau 64). Jenis siN yang ditandai merepresentasikan nilai bilangan bulat dari -2^(N-1) hingga 2^(N-1)-1 inklusif, dan jenis uiN yang tidak ditandai merepresentasikan nilai bilangan bulat dari 0 hingga 2^N-1 inklusif.
• Jenis floating point dapat berupa salah satu dari berikut:
  • Bilangan floating point 8-bit f8E3M4, f8E4M3, dan f8E5M2 yang mengikuti konvensi IEEE-754.
  • Jenis f8E4M3FN dan f8E5M2 masing-masing sesuai dengan encoding E4M3 dan E5M2 dari format FP8 yang dijelaskan dalam FP8 Formats for Deep Learning.
  • Jenis f8E4M3FNUZ dan f8E5M2FNUZ yang sesuai dengan encoding E4M3 dan E5M2 format FP8 yang dijelaskan dalam Format Numerik 8-bit untuk Jaringan Neural Dalam (Deep Neural Network).
  • Jenis f8E4M3B11FNUZ yang sesuai dengan encoding E4M3 format FP8 yang dijelaskan dalam Pelatihan dan Inferensi Floating Point 8-bit Hibrida (HFP8) untuk Jaringan Neural Dalam.
  • Jenis bf16 yang sesuai dengan format bfloat16 yang dijelaskan dalam BFloat16: Rahasia performa tinggi di Cloud TPU.
  • Jenis f16, f32, dan f64 yang masing-masing sesuai dengan binary16 ("presisi setengah"), binary32 ("presisi tunggal"), dan binary64 ("presisi ganda") yang dijelaskan dalam standar IEEE 754.
  • Jenis tf32 sesuai dengan format TensorFloat32 dan memiliki dukungan terbatas di StableHLO.
  • f4E2M1FN, f6E2M3FN, f6E3M2FN, dan f8E8M0FNU jenis MX (penskalaan mikro) yang dijelaskan dalam Spesifikasi Format Penskalaan Mikro OCP.
• Jenis kompleks merepresentasikan nilai kompleks yang memiliki bagian riil dan bagian imajiner dari jenis elemen yang sama. Jenis kompleks yang didukung adalah complex<f32> (kedua bagian berjenis f32) dan complex<f64> (kedua bagian berjenis f64).

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

Jenis fungsi merepresentasikan fungsi bernama dan anonim. Mereka memiliki jenis input (daftar jenis di sisi kiri ->) dan jenis output (daftar jenis di sisi kanan ->). Dalam banyak bahasa pemrograman, jenis fungsi adalah kelas utama, tetapi tidak di StableHLO.

StringType ::= 'string'

Jenis string merepresentasikan urutan byte. Tidak seperti di banyak bahasa pemrograman, jenis string bukan kelas utama di StableHLO dan hanya digunakan untuk menentukan metadata statis untuk elemen program.

Operasi

Operasi StableHLO (yang juga disebut ops) merepresentasikan sekumpulan operasi tingkat tinggi yang tertutup dalam model machine learning. Seperti yang dibahas di atas, sintaksis StableHLO sangat terinspirasi oleh MLIR, yang belum tentu merupakan alternatif paling ergonomis, tetapi bisa dibilang paling cocok untuk tujuan StableHLO dalam menciptakan lebih banyak interoperabilitas antara framework ML dan compiler ML.

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

Operasi StableHLO (yang juga disebut ops) memiliki nama, input/output, dan tanda tangan. Nama terdiri dari awalan stablehlo. dan mnemonik yang secara unik mengidentifikasi salah satu operasi yang didukung. Lihat di bawah untuk mengetahui daftar lengkap semua operasi yang didukung.

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

Operasi menggunakan input dan menghasilkan output. Input dikategorikan ke dalam nilai input (dihitung selama eksekusi), fungsi input (disediakan secara statis, karena dalam fungsi StableHLO bukan nilai kelas satu) dan atribut input (juga disediakan secara statis). Jenis input dan output yang digunakan dan dihasilkan oleh operasi bergantung pada mnemoniknya. Misalnya, operasi add menggunakan 2 nilai input dan menghasilkan 1 nilai output. Sebagai perbandingan, operasi select_and_scatter menggunakan 3 nilai input, 2 fungsi input, dan 3 atribut input.

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

Fungsi input (yang juga disebut fungsi anonim) sangat mirip dengan fungsi bernama, kecuali: 1) fungsi ini tidak memiliki ID (oleh karena itu disebut "anonim"), 2) fungsi ini tidak mendeklarasikan jenis output (jenis output disimpulkan dari operasi return dalam fungsi).

Sintaksis untuk fungsi input mencakup bagian yang saat ini tidak digunakan (lihat produksi Unused di atas) yang ada untuk kompatibilitas dengan MLIR. Di MLIR, ada konsep "region" yang lebih umum yang dapat memiliki beberapa "blok" operasi yang terhubung bersama melalui operasi lompatan. Blok ini memiliki ID yang sesuai dengan produksi Unused, sehingga dapat dibedakan satu sama lain. StableHLO tidak memiliki operasi lompatan, sehingga bagian sintaksis MLIR yang sesuai tidak digunakan (tetapi masih ada).

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

Atribut input memiliki nama dan nilai yang merupakan salah satu konstanta yang didukung. Ini adalah cara utama untuk menentukan metadata statis untuk elemen program. Misalnya, operasi concatenate menggunakan atribut dimension untuk menentukan dimensi yang digunakan untuk menggabungkan nilai inputnya. Demikian pula, operasi slice menggunakan beberapa atribut seperti start_indices dan limit_indices untuk menentukan batas yang digunakan untuk mengiris nilai input.

Saat ini, program StableHLO di luar sana terkadang berisi atribut yang tidak dijelaskan dalam dokumen ini. Ke depannya, kami berencana untuk menggabungkan atribut ini ke dalam opset StableHLO atau melarangnya muncul dalam program StableHLO. Sementara itu, berikut daftar atribut tersebut:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Metadata lokasi (#594).

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

Tanda tangan op terdiri dari jenis semua nilai input (daftar jenis di sisi kiri ->) dan jenis semua nilai output (daftar jenis di sisi kanan ->). Sebenarnya, jenis input bersifat redundan, dan jenis output juga hampir selalu redundan (karena untuk sebagian besar operasi StableHLO, jenis output dapat disimpulkan dari input). Namun demikian, tanda tangan op sengaja menjadi bagian dari sintaksis StableHLO untuk kompatibilitas dengan MLIR.

Berikut adalah contoh operasi yang mnemoniknya adalah select_and_scatter. Fungsi ini menggunakan 3 nilai input (%operand, %source, dan %init_value), 2 fungsi input , dan 3 atribut input (window_dimensions, window_strides, dan padding). Perhatikan cara tanda tangan op hanya menyertakan jenis nilai inputnya (tetapi bukan jenis fungsi dan atribut input yang disediakan secara inline).

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

Konstanta

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

Konstanta StableHLO memiliki literal dan jenis yang bersama-sama merepresentasikan nilai StableHLO. Umumnya, jenis adalah bagian dari sintaksis konstanta, kecuali jika tidak ambigu (misalnya, konstanta boolean tidak ambigu memiliki jenis i1, sedangkan konstanta bilangan bulat dapat memiliki beberapa kemungkinan jenis).

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

Konstanta boolean merepresentasikan nilai boolean true dan false. Konstanta Boolean memiliki jenis i1.

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

Konstanta bilangan bulat merepresentasikan nilai bilangan bulat melalui string yang menggunakan notasi desimal atau heksadesimal. Basis lainnya, misalnya biner atau oktal, tidak didukung. Konstanta bilangan bulat memiliki batasan berikut:

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

Konstanta floating point merepresentasikan nilai floating point melalui string yang menggunakan notasi desimal atau ilmiah. Selain itu, notasi heksadesimal dapat digunakan untuk menentukan secara langsung bit yang mendasarinya dalam format floating point dari jenis yang sesuai. Konstanta floating point memiliki batasan berikut:

• (C1) Jika notasi non-heksadesimal digunakan, is_wellformed(float_literal, float_type).
• (C2) Jika notasi heksadesimal digunakan, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Konstanta kompleks merepresentasikan nilai kompleks menggunakan daftar bagian real (muncul pertama) dan bagian imajiner (muncul kedua). Misalnya, (1.0, 0.0) : complex<f32> mewakili 1.0 + 0.0i, dan (0.0, 1.0) : complex<f32> mewakili 0.0 + 1.0i. Urutan bagian-bagian ini disimpan dalam memori ditentukan oleh implementasi. Konstanta kompleks memiliki batasan berikut:

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

Konstanta tensor merepresentasikan nilai tensor menggunakan daftar bertingkat yang ditentukan melalui notasi NumPy. Misalnya, dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> merepresentasikan nilai tensor dengan pemetaan berikut dari indeks ke elemen: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Urutan elemen ini disimpan dalam memori ditentukan oleh implementasi. Konstanta tensor memiliki batasan berikut:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), dengan:
  • 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)), dengan:
  • 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:]).
  • jika tidak, false.

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

Konstanta tensor terkuantisasi merepresentasikan nilai tensor terkuantisasi menggunakan notasi yang sama dengan konstanta tensor, dengan elemen yang ditentukan sebagai konstanta jenis penyimpanan. Konstanta tensor terkuantisasi memiliki batasan berikut:

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

Literal string terdiri dari byte yang ditentukan menggunakan karakter ASCII dan urutan escape. Byte ini tidak bergantung pada encoding, sehingga interpretasi byte ini ditentukan oleh implementasi. Literal string memiliki jenis string.

Operasi

abs

Semantik

Melakukan operasi abs per elemen pada tensor operand dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk bilangan bulat bertanda: modulus bilangan bulat.
• Untuk float: abs dari IEEE-754.
• Untuk bilangan kompleks: modulus kompleks.
• Untuk jenis yang dikuantisasi: dequantize_op_quantize(abs, operand, type(result)).

Input

Output

Batasan

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) ditentukan sebagai:
  • complex_element_type(element_type(operand)) if is_complex(operand).
  • baseline_element_type(operand) jika tidak.

Contoh

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

 Contoh Lainnya

tambahkan

Semantik

Melakukan penambahan per elemen dari dua tensor lhs dan rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk boolean: OR logis.
• Untuk bilangan bulat: penambahan bilangan bulat.
• Untuk float: addition dari IEEE-754.
• Untuk bilangan kompleks: penjumlahan kompleks.
• Untuk jenis yang dikuantisasi: dequantize_op_quantize(add, lhs, rhs, type(result)).

Input

Output

Batasan

• Jika operasi menggunakan tensor yang tidak dikuantisasi:
  • (C1) type(lhs) = type(rhs) = type(result).
• Jika operasi menggunakan tensor terkuantisasi:
  • (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) Jika is_per_axis_quantized(lhs), maka quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Jika is_per_axis_quantized(rhs), maka quantization_dimension(rhs) = quantization_dimension(result).

Contoh

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

 Contoh Lainnya

after_all

Semantik

Memastikan bahwa operasi yang menghasilkan inputs dijalankan sebelum operasi apa pun yang bergantung pada result. Eksekusi operasi ini tidak melakukan apa pun, operasi ini hanya ada untuk membuat dependensi data dari result ke inputs.

Input

Output

Contoh

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

 Contoh Lainnya

all_gather

Semantik

Dalam setiap grup proses di petak proses StableHLO, menggabungkan nilai tensor operands dari setiap proses di sepanjang all_gather_dim dan menghasilkan tensor results.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

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

Setelah itu, di dalam setiap process_group:

• operands...@receiver = [operand@sender for sender in process_group] untuk semua receiver di process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) untuk semua process di process_group.

Input

Output

Batasan

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_replicas jika cross_replica_and_partition digunakan.
  • num_processes jika flattened_ids digunakan.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Jika use_global_device_ids = true, maka channel_id > 0.
• (C6) type(results...) = type(operands...) kecuali:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Contoh

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

 Contoh Lainnya

all_reduce

Semantik

Dalam setiap grup proses di petak proses StableHLO, menerapkan fungsi reduksi computation ke nilai tensor operands dari setiap proses dan menghasilkan tensor results.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

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

Setelah itu, di dalam setiap process_group:

• results...@process[result_index] = exec(schedule) untuk beberapa pohon biner schedule dengan:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule adalah pohon biner yang ditentukan implementasinya yang traversal in-order-nya adalah to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Input

Output

Batasan

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_replicas jika cross_replica_and_partition digunakan.
  • num_processes jika flattened_ids digunakan.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Jika use_global_device_ids = true, maka channel_id > 0.
• (C5) computation memiliki jenis (tensor<E>, tensor<E>) -> (tensor<E>) dengan is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Contoh

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

 Contoh Lainnya

all_to_all

Semantik

Dalam setiap grup proses di petak proses StableHLO, membagi nilai tensor operands di sepanjang split_dimension menjadi beberapa bagian, menyebarkan bagian yang dibagi di antara proses, menggabungkan bagian yang disebarkan di sepanjang concat_dimension, dan menghasilkan tensor results. Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

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

Setelah itu, di dalam setiap process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) untuk semua sender di process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] di mana receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Input

Output

Batasan

• (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) ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_partitions jika cross_partition digunakan.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) kecuali, jika split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Contoh

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

 Contoh Lainnya

dan

Semantik

Melakukan AND per elemen dari dua tensor lhs dan rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk boolean: AND logis.
• Untuk bilangan bulat: bitwise AND.

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

atan2

Semantik

Melakukan operasi atan2 per elemen pada tensor lhs dan rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk float: atan2 dari IEEE-754.
• Untuk bilangan kompleks: atan2 kompleks.
• Untuk jenis yang dikuantisasi: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

batch_norm_grad

Semantik

Menghitung gradien beberapa input batch_norm_training dengan backpropagation dari grad_output, dan menghasilkan tensor grad_operand, grad_scale, dan grad_offset. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO yang ada menggunakan sintaksis Python sebagai berikut:

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

Untuk jenis yang dikuantisasi, melakukan 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)).

Input

Output

Batasan

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale, dan grad_offset memiliki baseline_element_type yang sama.
• (C3) operand, grad_output, dan grad_operand memiliki bentuk yang sama.
• (C4) scale, mean, variance, grad_scale, dan grad_offset memiliki bentuk yang sama.
• (C5) size(scale) = dim(operand, feature_index).

Contoh

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
// %grad_output: [
//                [[0.1, 0.1], [0.1, 0.1]],
//                [[0.1, 0.1], [0.1, 0.1]]
//               ]
%grad_operand, %grad_scale, %grad_offset =
"stablehlo.batch_norm_grad"(%operand, %scale, %mean, %variance, %grad_output) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>,
     tensor<2x2x2xf64>) -> (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %grad_operand: [
//                 [[0.0, 0.0], [0.0, 0.0]],
//                 [[0.0, 0.0], [0.0, 0.0]]
//                ]
// %grad_scale:  [0.0, 0.0]
// %grad_offset: [0.4, 0.4]

batch_norm_inference

Semantik

Menormalisasi tensor operand di semua dimensi kecuali dimensi feature_index dan menghasilkan tensor result. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO yang ada menggunakan sintaksis Python sebagai berikut:

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)

Untuk jenis yang dikuantisasi, melakukan 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)).

Input

Output

Batasan

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance, dan result memiliki baseline_element_type yang sama.
• (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).

Contoh

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

Semantik

Menghitung rata-rata dan varians di semua dimensi kecuali dimensi feature_index dan menormalisasi tensor operand yang menghasilkan tensor output, batch_mean dan batch_var. Secara lebih formal, operasi ini dapat dinyatakan sebagai dekomposisi ke operasi StableHLO yang ada menggunakan sintaksis Python sebagai berikut:

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

Untuk jenis yang dikuantisasi, melakukan 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)).

Input

Output

Batasan

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var, dan output memiliki baseline_element_type yang sama.
• (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).

Contoh

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

Semantik

Melakukan operasi bitcast pada tensor operand dan menghasilkan tensor result di mana bit seluruh tensor operand ditafsirkan ulang menggunakan jenis tensor result.

Secara lebih formal, dengan kondisi E = element_type(operand), E' = element_type(result), dan R = rank(operand):

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

bits menampilkan representasi dalam memori dari nilai tertentu, dan perilakunya ditetapkan oleh implementasi karena representasi tensor yang tepat ditetapkan oleh implementasi, dan representasi jenis elemen yang tepat juga ditetapkan oleh implementasi.

Input

Output

Batasan

• (C1) Dengan E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result), dan R = rank(operand):
  • Jika num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Jika num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) untuk semua 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Jika num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) untuk semua 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Jika is_complex(operand) or is_complex(result), maka is_complex(operand) and is_complex(result).

Contoh

// %operand: 0x0123456789ABCDEF
%result = "stablehlo.bitcast_convert"(%operand) : (tensor<f64>) -> tensor<4xf16>
// %result: [0xCDEF, 0x89AB, 0x4567, 0x0123] // little-endian representation

 Contoh Lainnya

broadcast_in_dim

Semantik

Memperluas dimensi dan/atau peringkat tensor input dengan menduplikasi data dalam tensor operand dan menghasilkan tensor result. Secara lebih formal, result[result_index] = operand[operand_index] dengan semua d dalam axes(operand):

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

Input

Output

Batasan

• (C1) element_type(result) diberikan oleh:
  • element_type(operand), jika !is_per_axis_quantized(operand).
  • element_type(operand), kecuali quantization_dimension(operand), scales(operand), dan zero_points(operand) dapat berbeda dari quantization_dimension(result), scales(result), dan zero_points(result) masing-masing, jika tidak.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Untuk semua d di axes(operand):
  • dim(operand, d) = 1 atau
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Jika is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Jika dim(operand, quantization_dimension(operand)) = 1, maka scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Contoh

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

 Contoh Lainnya

casing

Semantik

Menghasilkan output dari eksekusi tepat satu fungsi dari branches bergantung pada nilai index. Secara lebih formal, result = selected_branch() dengan:

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

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

cbrt

Semantik

Melakukan operasi akar kubik per elemen pada tensor operand dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk float: rootn(x, 3) dari IEEE-754.
• Untuk bilangan kompleks: akar kubik kompleks.
• Untuk jenis yang dikuantisasi: dequantize_op_quantize(cbrt, operand, type(result))

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

ceil

Semantik

Melakukan ceil per elemen pada tensor operand dan menghasilkan tensor result. Menerapkan operasi roundToIntegralTowardPositive dari spesifikasi IEEE-754. Untuk jenis yang dikuantisasi, melakukan dequantize_op_quantize(ceil, operand, type(result)).

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

cholesky

Semantik

Menghitung dekomposisi Cholesky dari sekumpulan matriks.

Secara lebih formal, untuk semua i dalam index_space(result), result[i0, ..., iR-3, :, :] adalah dekomposisi Cholesky dari a[i0, ..., iR-3, :, :], dalam bentuk matriks segitiga bawah (jika lower adalah true) atau segitiga atas (jika lower adalah false). Nilai output dalam segitiga yang berlawanan, yaitu segitiga atas ketat atau segitiga bawah ketat, ditentukan oleh implementasi.

Jika ada i di mana matriks input bukan matriks definit positif Hermitian, maka perilakunya tidak ditentukan.

Untuk jenis yang dikuantisasi, melakukan dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Input

Output

Batasan

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

Contoh

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

penjepit

Semantik

Mengepaskan setiap elemen tensor operand antara nilai minimum dan maksimum serta menghasilkan tensor result. Secara lebih formal, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), dengan min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Untuk jenis yang dikuantisasi, melakukan dequantize_op_quantize(clamp, min, operand, max, type(result)).

Menerapkan pengurutan pada bilangan kompleks melibatkan semantik yang mengejutkan, jadi pada masa mendatang kami berencana menghapus dukungan untuk bilangan kompleks untuk operasi ini (#560).

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

collective_broadcast

Semantik

Dalam setiap grup proses di petak proses StableHLO, kirim nilai tensor operand dari proses sumber ke proses target dan hasilkan tensor result.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

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

Setelah itu, result@process diberikan oleh:

• operand@process_groups[i, 0] jika ada i sehingga prosesnya berada di process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) sebaliknya.

Input

Output

Batasan

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N dengan N ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_partitions jika cross_partition digunakan.
• (C3) type(result) = type(operand).

Contoh

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

Semantik

Dalam setiap grup proses di petak proses StableHLO, mengirimkan nilai tensor operand dari proses sumber ke proses target dan menghasilkan tensor result.

Operasi ini membagi petak proses StableHLO menjadi process_groups yang ditentukan sebagai berikut:

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

Setelah itu, result@process diberikan oleh:

• operand@process_groups[i, 0], jika ada i sehingga process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) sebaliknya.

Input

Output

Batasan

• (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, dengan N ditentukan sebagai:
  • num_replicas jika cross_replica digunakan.
  • num_partitions jika cross_partition digunakan.
• (C5) type(result) = type(operand).

Contoh

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

 Contoh Lainnya

bandingkan

Semantik

Melakukan perbandingan elemen demi elemen pada tensor lhs dan rhs sesuai dengan comparison_direction dan compare_type, serta menghasilkan tensor result.

Nilai comparison_direction dan compare_type memiliki semantik berikut:

Untuk jenis elemen boolean dan bilangan bulat:

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

Untuk jenis elemen floating point dengan compare_type = FLOAT, op mengimplementasikan operasi IEEE-754 berikut:

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

Untuk jenis elemen floating point dengan compare_type = TOTALORDER, op menggunakan kombinasi operasi totalOrder dan compareQuietEqual dari IEEE-754.

Untuk jenis elemen yang kompleks, perbandingan leksikografis pasangan (real, imag) dilakukan menggunakan comparison_direction dan compare_type yang diberikan. Pemberlakuan pengurutan pada bilangan kompleks melibatkan semantik yang mengejutkan, jadi pada masa mendatang kami berencana menghapus dukungan untuk bilangan kompleks jika comparison_direction adalah GE, GT, LE, atau LT (#560).

Untuk jenis yang dikuantisasi. Melakukan dequantize_compare(lhs, rhs, comparison_direction).

Input

Output

Batasan

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type ditentukan sebagai:
  • SIGNED if is_signed_integer(element_type(lhs)).
  • UNSIGNED if is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT atau TOTALORDER jika is_float(element_type(lhs)).
  • FLOAT if is_complex(element_type(lhs)).

Contoh

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

 Contoh Lainnya

kompleks

Semantik

Melakukan konversi per elemen ke nilai kompleks dari pasangan nilai riil dan imajiner, lhs dan rhs, serta menghasilkan tensor result.

Input

Output

Batasan

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) memiliki jenis complex<E> dengan E = element_type(lhs).

Contoh

// %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)]

 Contoh Lainnya

gabungan

Semantik

Mengkapsulasi operasi yang terdiri dari (dikomposisikan) operasi StableHLO lainnya, mengambil inputs dan composite_attributes serta menghasilkan results. Semantik op diimplementasikan oleh atribut decomposition. Operasi composite dapat diganti dengan dekomposisinya tanpa mengubah semantik program. Jika penyusunan sebaris dekomposisi tidak memberikan semantik op yang sama, sebaiknya gunakan custom_call.

Kolom version (defaultnya 0) digunakan untuk menunjukkan kapan semantik komposit berubah.

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

concatenate

Semantik

Menggabungkan inputs di sepanjang dimensi dimension dalam urutan yang sama seperti argumen yang diberikan dan menghasilkan tensor result. Secara lebih formal, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], dengan:

1. id = d0 + ... + dk-1 + kd.
2. d sama dengan dimension, dan d0, ... adalah ukuran dimensi ke-d dari inputs.

Input

Output

Batasan

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) kecuali 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]) kecuali:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Contoh

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

 Contoh Lainnya

konstanta

Semantik

Menghasilkan tensor output dari konstanta value.

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

melakukan konversi

Semantik

Melakukan konversi per elemen dari satu jenis elemen ke jenis elemen lainnya pada tensor operand dan menghasilkan tensor result.

Untuk konversi boolean-to-any-supported-type, nilai false dikonversi menjadi nol, dan nilai true dikonversi menjadi satu. Untuk konversi any-supported-type-to-boolean, nilai nol dikonversi menjadi false, dan nilai bukan nol dikonversi menjadi true. Lihat di bawah untuk mengetahui cara kerja ini untuk jenis yang kompleks.

Untuk konversi yang melibatkan bilangan bulat ke bilangan bulat, bilangan bulat ke floating point atau floating point ke floating point, jika nilai sumber dapat direpresentasikan secara persis dalam jenis tujuan, nilai hasilnya adalah representasi persis tersebut. Jika tidak, perilakunya akan ditentukan kemudian (#180).

Untuk konversi yang melibatkan floating-point-to-integer, bagian pecahan akan dipangkas. Jika nilai yang dipangkas tidak dapat ditampilkan dalam jenis tujuan, perilakunya akan ditentukan (#180).

Konversi yang melibatkan kompleks ke kompleks mengikuti perilaku yang sama dengan konversi floating point ke floating point untuk mengonversi bagian riil dan imajiner.

Untuk konversi complex-to-any-other-type dan any-other-type-to-complex, nilai imajiner sumber diabaikan atau nilai imajiner tujuan dibuat nol. Konversi bagian nyata mengikuti konversi floating point.

Pada prinsipnya, operasi ini dapat mengekspresikan dekuantisasi (konversi dari tensor terkuantisasi ke tensor reguler), kuantisasi (konversi dari tensor reguler ke tensor terkuantisasi), dan rekuantisasi (konversi antara tensor terkuantisasi), tetapi saat ini kita memiliki operasi khusus untuk itu - uniform_dequantize untuk kasus penggunaan pertama dan uniform_quantize untuk kasus penggunaan kedua dan ketiga. Pada masa mendatang, kedua operasi ini dapat digabungkan menjadi convert (#1576).

Input

Output

Batasan

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

Contoh

// %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)]

 Contoh Lainnya

konvolusi

Semantik

Menghitung perkalian titik antara jendela lhs dan irisan rhs serta menghasilkan result. Diagram berikut menunjukkan cara elemen di result dihitung dari lhs dan rhs menggunakan contoh konkret.

Secara lebih formal, pertimbangkan pembingkaian ulang input berikut dalam hal lhs agar dapat mengekspresikan jendela 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).

Pembingkaian ulang ini menggunakan fungsi bantuan berikut:

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

Jika feature_group_count = 1 dan batch_group_count = 1, maka untuk semua output_spatial_index di index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product dengan:

• 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]). Fitur ini tampaknya tidak digunakan, jadi pada masa mendatang kami berencana menghapusnya (#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]).

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

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

Untuk jenis yang dikuantisasi, melakukan 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)).

Untuk jenis terkuantisasi hibrida, melakukan 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).

Input

Output

Batasan

• (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) Diberikan 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) Mengingat 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) Dengan 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) didefinisikan sebagai:
  • 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 jika tidak, di mana:
  • 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.
• Jika operasi menggunakan tensor yang tidak dikuantisasi:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Jika operasi menggunakan tensor terkuantisasi:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Jika is_per_axis_quantized(rhs), maka quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Jika is_per_axis_quantized(result), maka quantization_dimension(result) = output_feature_dimension.
  • Jika is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Jika is_per_tensor_quantized(rhs), maka is_per_tensor_quantized(result).
  • Jika !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Contoh

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

 Contoh Lainnya

kosinus

Semantik

Melakukan operasi kosinus per elemen pada tensor operand dan menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk float: cos dari IEEE-754.
• Untuk bilangan kompleks: kosinus kompleks.
• Untuk jenis yang dikuantisasi: dequantize_op_quantize(cosine, operand, type(result)).

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

count_leading_zeros

Semantik

Melakukan penghitungan per elemen dari jumlah bit nol di awal dalam tensor operand dan menghasilkan tensor result.

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

custom_call

Semantik

Mengkapsulasi operasi call_target_name yang ditentukan implementasinya yang mengambil inputs dan called_computations serta menghasilkan results. has_side_effect, backend_config, dan api_version dapat digunakan untuk memberikan metadata tambahan yang ditentukan implementasi.

Saat ini, operasi ini berisi kumpulan metadata yang cukup tidak teratur yang mencerminkan evolusi organik operasi yang setara di compiler XLA. Pada masa mendatang, kami berencana menyatukan metadata ini (#741).

Input

Output

(Dukungan GPU XLA) Target custom_call khusus

Ada tiga call_target_name khusus yang terkait dengan jenis buffer: CreateBuffer membuat buffer yang belum diinisialisasi, Pin membuat buffer yang sudah diinisialisasi, dan Unpin membatalkan alokasi buffer serta menampilkan konten 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>

Alias

Beberapa operasi custom_call mungkin memerlukan bagian dalam output dan bagian dalam operan untuk berbagi memori yang sama. Hal ini dapat dinyatakan melalui output_operand_aliases. Representasi pasangan alias terdiri dari daftar indeks tuple output yang merepresentasikan bagian output, dan operand_index beserta daftar indeks tuple operand yang merepresentasikan bagian operand. Daftar output atau indeks tuple operand kosong jika jenis yang sesuai bukan jenis tuple, dan dapat memiliki panjang yang berubah-ubah untuk jenis tuple yang bersarang secara berubah-ubah. Hal ini mirip dengan representasi alias XLA.

Bagian output dan bagian input dalam pasangan alias harus memiliki jenis yang sama. Untuk operasi custom_call yang bukan panggilan ke CreateBuffer, Pin, dan Unpin, operand buffer dapat muncul di paling banyak satu pasangan alias, dan output buffer harus muncul di satu pasangan alias.

Contoh

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

bagi

Semantik

Melakukan pembagian per elemen tensor dividen lhs dan tensor pembagi rhs serta menghasilkan tensor result. Bergantung pada jenis elemen, lakukan tindakan berikut:

• Untuk bilangan bulat: pembagian bilangan bulat yang menghasilkan hasil bagi aljabar dengan bagian pecahan yang diabaikan.
• Untuk float: division dari IEEE-754.
• Untuk bilangan kompleks: pembagian kompleks.
• Untuk jenis yang dikuantisasi:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Input

Output

Batasan

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

Contoh

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

 Contoh Lainnya

dot_general

Semantik

Menghitung perkalian titik antara slice lhs dan slice rhs serta menghasilkan tensor result.

Secara lebih formal, result[result_index] = dot_product, dengan:

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

Untuk jenis yang dikuantisasi, melakukan 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)).

Untuk jenis terkuantisasi hibrida, melakukan 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 mengontrol kompromi antara kecepatan dan akurasi untuk komputasi di backend akselerator. Ini dapat berupa salah satu dari nilai berikut (saat ini, semantik nilai enum ini kurang ditentukan, tetapi kami berencana untuk menanganinya di #755):

• DEFAULT: Perhitungan tercepat, tetapi perkiraan paling tidak akurat untuk angka asli.
• HIGH: Perhitungan lebih lambat, tetapi perkiraan yang lebih akurat untuk angka asli.
• HIGHEST: Perhitungan paling lambat, tetapi perkiraan paling akurat untuk angka asli.

DotAlgorithm menentukan properti utama algoritma yang digunakan untuk menerapkan operasi titik, yang juga menentukan presisi. Jika kolom atribut algoritma disetel, maka precision_config harus DEFAULT. DotAlgorithms tidak memiliki nilai default, karena parameter default ditentukan oleh penerapan. Dengan demikian, semua kolom algoritma titik dapat disetel ke None untuk menentukan algoritma titik kosong, yang akan menggunakan nilai precision_config.

Kolom DotAlgorithm mencakup:

• lhs_precision_type dan rhs_precision_type, presisi yang dibulatkan untuk LHS dan RHS operasi. Jenis presisi tidak bergantung pada jenis penyimpanan input dan output.
• accumulation_type presisi yang digunakan untuk akumulasi.
• lhs_component_count, rhs_component_count, dan num_primitive_operations berlaku saat kita melakukan algoritma yang menguraikan LHS dan/atau RHS menjadi beberapa komponen dan melakukan beberapa operasi titik "primitif" pada nilai-nilai tersebut - biasanya untuk mengemulasi presisi yang lebih tinggi (misalnya, Memanfaatkan Jenis Data Kecerdasan Buatan bfloat16 untuk Komputasi Presisi Tinggi: bf16_6x tf32_3x, dll.). Untuk algoritma tanpa dekomposisi, nilai ini harus ditetapkan ke 1.
• allow_imprecise_accumulation untuk menentukan apakah akumulasi dalam presisi yang lebih rendah diizinkan untuk beberapa langkah (misalnya, CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Contoh atribut 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}

Implementasi yang akan memutuskan kombinasi mana yang didukung. Secara umum, tidak ada jaminan bahwa setiap algoritma didukung di setiap jenis akselerator oleh konsumen StableHLO. Jika algoritma tertentu tidak didukung, error akan muncul, bukan beralih ke alternatif. Verifikasi StableHLO akan memberikan verifikasi upaya terbaik, mencegah algoritma yang tidak diketahui didukung di hardware mana pun.

Lihat xla_data.proto > Algorithm untuk mengetahui beberapa nilai algoritma yang didukung. Tiket #2483 mencatat rencana untuk membuat dokumen terpusat tentang algoritma yang didukung oleh backend.

Input

Output

Batasan

• (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).
• Jika operasi menggunakan tensor yang tidak dikuantisasi:
  • (C13) element_type(lhs) = element_type(rhs).
• Jika operasi menggunakan tensor terkuantisasi:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Jika is_per_axis_quantized(rhs), maka quantization_dimension(rhs) tidak ada di rhs_contracting_dimensions.
  • Jika is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Jika is_per_tensor_quantized(rhs), maka is_per_tensor_quantized(result).
  • Jika !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Jika !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.

Contoh

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

 Contoh Lainnya

dynamic_broadcast_in_dim

Semantik

Operasi ini secara fungsional identik dengan broadcast_in_dim op, tetapi bentuk hasilnya ditentukan secara dinamis melalui output_dimensions.

Operasi ini juga menerima atribut opsional known_expanding_dimensions, known_nonexpanding_dimensions untuk menyatakan pengetahuan statis tentang perilaku perluasan dimensi. Jika tidak ditentukan, semua dimensi diasumsikan dapat diperluas.

Input

Output

Batasan

• (C1) element_type(result) diberikan oleh:
  • element_type(operand), jika !is_per_axis_quantized(operand).
  • element_type(operand), kecuali quantization_dimension(operand), scales(operand), dan zero_points(operand) dapat berbeda dari quantization_dimension(result), scales(result), dan zero_points(result) masing-masing, jika tidak.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Untuk semua d di axes(operand):
  • dim(operand, d) = 1 atau
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Jika is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Jika dim(operand, quantization_dimension(operand)) = 1, maka 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).

Contoh

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

 Contoh Lainnya

dynamic_conv

Semantik

Operasi ini secara fungsional identik dengan operasi konvolusi, tetapi padding ditentukan secara dinamis melalui padding.

Input