Thông số kỹ thuật của StableHLO

StableHLO là một tập hợp thao tác dành cho các thao tác cấp cao (HLO) trong mô hình học máy (ML). StableHLO hoạt động như một lớp khả năng di chuyển giữa các khung ML và trình biên dịch ML khác nhau: các khung ML tạo ra chương trình StableHLO tương thích với các trình biên dịch ML sử dụng chương trình StableHLO.

Mục tiêu của chúng tôi là đơn giản hoá và đẩy nhanh quá trình phát triển công nghệ học máy bằng cách tăng khả năng tương tác giữa nhiều khung học máy (chẳng hạn như TensorFlow, JAX và PyTorch) và các trình biên dịch học máy (như XLA và IREE). Để đạt được mục tiêu đó, tài liệu này cung cấp thông số kỹ thuật cho ngôn ngữ lập trình StableHLO.

Quy cách này bao gồm 3 phần chính. Trước tiên, phần Chương trình mô tả cấu trúc của các chương trình ổn định HLO, trong đó bao gồm các hàm ổn định HLO, trong đó có các hoạt động ổn định HLO. Trong cấu trúc đó, phần Ops (Thao tác) chỉ định ngữ nghĩa của từng thao tác. Phần Thực thi cung cấp ngữ nghĩa cho tất cả các thao tác này thực thi cùng nhau trong một chương trình. Cuối cùng, phần Ký hiệu thảo luận về ký hiệu được sử dụng trong toàn bộ quy cách.

Để xem thông số kỹ thuật của bản phát hành StableHLO trước đây, hãy mở kho lưu trữ tại bản phát hành được gắn thẻ mà bạn quan tâm. Ví dụ: Thông số kỹ thuật StableHLO v0.19.0. Để xem các thay đổi xảy ra ở mỗi lần nâng cấp phiên bản nhỏ của StableHLO, hãy tham khảo nhật ký phiên bản trong VhloDialect.td.

Chương trình

Program ::= {Func}

Chương trình StableHLO bao gồm một số lượng hàm StableHLO tuỳ ý. Dưới đây là một chương trình mẫu có hàm @main có 3 đầu vào (%image, %weights và %bias) và 1 đầu ra. Phần nội dung của hàm có 6 toán tử.

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

Hàm

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

Hàm StableHLO (còn gọi là hàm được đặt tên) có một giá trị nhận dạng, dữ liệu đầu vào/đầu ra và phần nội dung. Trong tương lai, chúng tôi dự định giới thiệu thêm siêu dữ liệu cho các hàm để đạt được khả năng tương thích tốt hơn với HLO (#425, #626, #740, #744).

Giá trị nhận dạng

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

Giá trị nhận dạng ổn địnhHLO tương tự như giá trị nhận dạng trong nhiều ngôn ngữ lập trình, với hai điểm đặc biệt: 1) tất cả giá trị nhận dạng đều có các dấu hiệu phân biệt nhiều loại giá trị nhận dạng, 2) giá trị nhận dạng giá trị có thể hoàn toàn là số để đơn giản hoá việc tạo chương trình StableHLO.

Loại

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

Các loại StableHLO được phân loại thành các loại giá trị (còn gọi là các loại hạng nhất) đại diện cho các giá trị StableHLO và các loại không phải giá trị mô tả các phần tử chương trình khác. Các loại StableHLO tương tự như các loại trong nhiều ngôn ngữ lập trình, với đặc điểm chính là bản chất dành riêng cho miền của StableHLO, dẫn đến một số kết quả bất thường (ví dụ: các loại vô hướng không phải là loại giá trị).

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

Loại tensor biểu thị các tensor, tức là các mảng nhiều chiều. Các thành phần này có một hình dạng và một loại phần tử, trong đó hình dạng đại diện cho kích thước kích thước không âm hoặc không xác định theo thứ tự tăng dần của kích thước tương ứng (còn gọi là trục) được đánh số từ 0 đến R-1. Số lượng phương diện R được gọi là hạng. Ví dụ: tensor<2x3xf32> là một loại tensor có hình dạng 2x3 và loại phần tử f32. Mảng này có hai phương diện (hay nói cách khác là hai trục) – phương diện thứ 0 và phương diện thứ 1 – có kích thước là 2 và 3. Thứ hạng của nó là 2.

Các hình dạng có thể không xác định một phần hoặc hoàn toàn (động), ví dụ: tensor<?x2xf64> không xác định một phần và tensor<?x?xf64> không xác định hoàn toàn. Kích thước kích thước động được biểu thị bằng ?. Không thể xếp hạng hình dạng.

Trong tương lai, chúng tôi dự định mở rộng các loại tensor ngoài kích thước và loại phần tử, chẳng hạn như bao gồm cả bố cục (#629) và độ thưa thớt (#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

Các loại phần tử được lượng tử hoá biểu thị các giá trị số nguyên của một loại bộ nhớ trong khoảng từ storage_min đến storage_max (bao gồm) tương ứng với các giá trị dấu phẩy động của một loại đã biểu thị. Đối với một giá trị số nguyên nhất định i, giá trị dấu phẩy động tương ứng f có thể được tính là f = (i - zero_point) * scale, trong đó scale và zero_point được gọi là tham số lượng tử hoá. storage_min và storage_max là không bắt buộc trong ngữ pháp, nhưng có giá trị mặc định là min_value(storage_type) và max_value(storage_type) tương ứng. Các loại phần tử được lượng tử hoá có các điều kiện ràng buộc sau:

• (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) Nếu là is_empty(quantization_dimension), thì size(scales) = 1.
• (C11) 0 <= quantization_dimension.

Hiện tại, QuantizationScale là một hằng số dấu phẩy động, nhưng có rất nhiều người quan tâm đến các tỷ lệ dựa trên số nguyên, được biểu thị bằng hệ số và độ lệch. Chúng tôi dự định sẽ khám phá vấn đề này trong tương lai gần (#1404).

Chúng tôi đang thảo luận về ngữ nghĩa của QuantizationZeroPoint, bao gồm cả loại, giá trị và liệu có thể chỉ có một hoặc nhiều điểm 0 trong một loại tensor được lượng tử hoá hay không. Dựa trên kết quả của cuộc thảo luận này, thông số kỹ thuật xung quanh điểm 0 có thể thay đổi trong tương lai (#1405).

Một cuộc thảo luận khác đang diễn ra liên quan đến ngữ nghĩa của QuantizationStorageMin và QuantizationStorageMax để xác định xem có áp đặt bất kỳ quy tắc ràng buộc nào đối với các giá trị này và đối với các giá trị của tensor lượng tử hoá hay không (#1406).

Cuối cùng, chúng tôi dự định khám phá cách thể hiện các điểm không xác định và điểm 0, tương tự như cách chúng tôi dự định khám phá cách thể hiện kích thước phương diện không xác định (#1407).

Các loại tensor lượng tử hoá biểu thị các tensor có phần tử lượng tử hoá. Các tensor này giống hệt như tensor thông thường, ngoại trừ việc các phần tử của chúng có loại phần tử lượng tử hoá thay vì kiểu phần tử thông thường.

Trong tensor lượng tử hoá, quá trình lượng tử hoá có thể là trên mỗi tensor, nghĩa là có một scale và zero_point cho toàn bộ tensor hoặc có thể là trên mỗi trục, nghĩa là có nhiều scales và zero_points, một cặp trên mỗi lát cắt của một kích thước cụ thể quantization_dimension. Chính xác hơn, trong một tensor t có lượng tử hoá theo trục, có dim(t, quantization_dimension) lát cắt của quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], v.v. Tất cả các phần tử trong lát cắt thứ i sử dụng scales[i] và zero_points[i] làm tham số lượng tử hoá. Các loại tensor lượng tử hoá có các điều kiện ràng buộc sau:

• Đối với việc lượng tử hoá theo từng tensor:
  • Không có quy tắc ràng buộc nào khác.
• Đối với việc lượng tử hoá theo mỗi trục:
  • (C12) quantization_dimension < rank(self).
  • (C13) dim(self, quantization_dimension) = size(scales).

TokenType ::= 'token'

Loại mã thông báo đại diện cho mã thông báo, tức là các giá trị mờ do một số thao tác tạo ra và sử dụng. Mã thông báo được dùng để áp đặt thứ tự thực thi cho các thao tác như mô tả trong phần Thực thi.

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

Loại dữ liệu tuple đại diện cho các tuple, tức là các danh sách không đồng nhất. Bộ dữ liệu là một tính năng cũ chỉ tồn tại để tương thích với HLO. Trong HLO, các bộ dữ liệu được dùng để biểu thị dữ liệu đầu vào và đầu ra biến đổi. Trong StableHLO, đầu vào và đầu ra biến đổi được hỗ trợ nguyên bản và việc sử dụng duy nhất các bộ dữ liệu trong StableHLO là để biểu thị toàn diện ABI HLO, trong đó, ví dụ: T, tuple<T> và tuple<tuple<T>> có thể khác nhau đáng kể tuỳ thuộc vào cách triển khai cụ thể. Trong tương lai, chúng tôi dự định thực hiện các thay đổi đối với HLO ABI để có thể xoá các loại bộ dữ liệu từ 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'

Loại phần tử đại diện cho các phần tử của loại tensor. Không giống như trong nhiều ngôn ngữ lập trình, các loại này không phải là lớp đầu tiên trong StableHLO. Điều này có nghĩa là các chương trình StableHLO không thể trực tiếp biểu thị các giá trị của các loại này (do đó, bạn nên biểu thị các giá trị vô hướng thuộc loại T bằng các giá trị tensor 0 chiều thuộc loại tensor<T>).

• Loại boolean đại diện cho các giá trị boolean true và false.
• Loại số nguyên có thể có dấu (si) hoặc không dấu (ui) và có một trong các chiều rộng bit được hỗ trợ (2, 4, 8, 16, 32 hoặc 64). Loại siN đã ký biểu thị giá trị số nguyên từ -2^(N-1) đến 2^(N-1)-1, và loại uiN không dấu biểu thị giá trị số nguyên từ 0 đến 2^N-1.
• Loại dấu phẩy động có thể là một trong những loại sau:
  • f8E3M4, f8E4M3 và f8E5M2 số dấu phẩy động 8 bit tuân theo quy ước IEEE-754.
  • Các loại f8E4M3FN và f8E5M2 tương ứng với các phương thức mã hoá E4M3 và E5M2 của định dạng FP8 được mô tả trong phần Định dạng FP8 cho học sâu.
  • Các loại f8E4M3FNUZ và f8E5M2FNUZ tương ứng với mã hoá E4M3 và E5M2 của các định dạng FP8 được mô tả trong phần Định dạng số 8 bit cho Mạng nơron sâu.
  • Loại f8E4M3B11FNUZ tương ứng với mã hoá E4M3 của các định dạng FP8 được mô tả trong bài viết Đào tạo và suy luận theo điểm nổi kết hợp 8 bit (HFP8) cho mạng nơron sâu.
  • Loại bf16 tương ứng với định dạng bfloat16 được mô tả trong BFloat16: Bí quyết đạt hiệu suất cao trên Cloud TPU.
  • Các loại f16, f32 và f64 tương ứng với các định dạng binary16 ("độ chính xác đơn"), binary32 ("độ chính xác đơn") và binary64 ("độ chính xác kép") được mô tả trong tiêu chuẩn IEEE 754.
  • Loại tf32 tương ứng với định dạng TensorFloat32 và có hỗ trợ hạn chế trong StableHLO.
  • Các loại f4E2M1FN, f6E2M3FN, f6E3M2FN và f8E8M0FNU MX (chia tỷ lệ theo tỷ lệ) được mô tả trong bài viết Thông số kỹ thuật về định dạng vi mô OCP.
• Loại phức đại diện cho các giá trị phức có phần thực và phần ảo thuộc cùng một loại phần tử. Các loại phức tạp được hỗ trợ là complex<f32> (cả hai phần đều thuộc loại f32) và complex<f64> (cả hai phần đều thuộc loại f64).

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

Loại hàm đại diện cho cả hàm được đặt tên và hàm ẩn danh. Các hàm này có loại đầu vào (danh sách các loại ở bên trái ->) và loại đầu ra (danh sách các loại ở bên phải ->). Trong nhiều ngôn ngữ lập trình, các loại hàm là lớp đầu tiên, nhưng không phải trong StableHLO.

StringType ::= 'string'

Loại chuỗi biểu thị chuỗi byte. Không giống như trong nhiều ngôn ngữ lập trình, loại chuỗi không phải là lớp đầu tiên trong StableHLO và chỉ được dùng để chỉ định siêu dữ liệu tĩnh cho các phần tử chương trình.

Hoạt động tính toán

Toán tử StableHLO (còn gọi là toán tử) đại diện cho một tập hợp kín các toán tử cấp cao trong mô hình học máy. Như đã thảo luận ở trên, cú pháp StableHLO được lấy cảm hứng rất nhiều từ MLIR. Đây không nhất thiết là giải pháp thay thế phù hợp nhất về mặt công thái học, nhưng có thể là giải pháp phù hợp nhất với mục tiêu của StableHLO là tạo ra khả năng tương tác cao hơn giữa các khung máy học và trình biên dịch máy học.

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

Thao tác StableHLO (còn gọi là ops) có tên, đầu vào/đầu ra và chữ ký. Tên này bao gồm tiền tố stablehlo. và một ký tự ghi nhớ giúp xác định duy nhất một trong các hoạt động được hỗ trợ. Hãy xem phần bên dưới để biết danh sách đầy đủ tất cả hoạt động được hỗ trợ.

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

Các thao tác sử dụng dữ liệu đầu vào và tạo dữ liệu đầu ra. Các giá trị đầu vào được phân loại thành giá trị đầu vào (được tính trong quá trình thực thi), hàm đầu vào (được cung cấp tĩnh vì trong các hàm StableHLO không phải là giá trị hạng nhất) và các thuộc tính đầu vào (cũng được cung cấp theo phương thức tĩnh). Loại dữ liệu đầu vào và đầu ra mà một toán tử tiêu thụ và tạo ra phụ thuộc vào ký hiệu của toán tử đó. Ví dụ: toán tử add sử dụng 2 giá trị đầu vào và tạo ra 1 giá trị đầu ra. So sánh, toán tử select_and_scatter sử dụng 3 giá trị đầu vào, 2 hàm đầu vào và 3 thuộc tính đầu vào.

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

Hàm đầu vào (còn được gọi là hàm ẩn danh) rất giống với hàm được đặt tên, ngoại trừ: 1) chúng không có giá trị nhận dạng (do đó có tên "anonymous"), 2) chúng không khai báo loại đầu ra (loại đầu ra được suy ra từ hoạt động return trong hàm).

Cú pháp cho các hàm đầu vào bao gồm một phần hiện không được sử dụng (xem quá trình tạo Unused ở trên) để tương thích với MLIR. Trong MLIR, có một khái niệm chung hơn về "vùng" có thể có nhiều "khối" hoạt động được kết nối với nhau thông qua hoạt động nhảy. Các khối này có mã nhận dạng tương ứng với quá trình tạo Unused để có thể phân biệt với nhau. StableHLO không có hoạt động nhảy, vì vậy, phần tương ứng của cú pháp MLIR không được sử dụng (nhưng vẫn ở đó).

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

Thuộc tính đầu vào có tên và giá trị là một trong các hằng số được hỗ trợ. Đây là phương thức chính để chỉ định siêu dữ liệu tĩnh cho các phần tử của chương trình. Ví dụ: toán tử concatenate sử dụng thuộc tính dimension để chỉ định phương diện theo đó các giá trị đầu vào của toán tử này được nối với nhau. Tương tự, toán tử slice sử dụng nhiều thuộc tính như start_indices và limit_indices để chỉ định các giới hạn dùng để cắt giá trị đầu vào.

Hiện tại, các chương trình StableHLO trong thực tế đôi khi chứa các thuộc tính không được mô tả trong tài liệu này. Trong tương lai, chúng tôi dự định đưa các thuộc tính này vào cơ chế ổn định HLO hoặc cấm các thuộc tính này xuất hiện trong các chương trình ổn định HLO. Trong thời gian chờ đợi, sau đây là danh sách các thuộc tính này:

• layout (#629).
• mhlo.frontend_attributes (#628).
• mhlo.sharding (#619).
• output_operand_aliases (#740).
• Siêu dữ liệu vị trí (#594).

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

Chữ ký toán tử bao gồm các loại của tất cả giá trị đầu vào (danh sách các loại ở bên trái của ->) và các loại của tất cả giá trị đầu ra (danh sách các loại ở bên phải của ->). Nói một cách nghiêm ngặt, các loại đầu vào là thừa và các loại đầu ra hầu như luôn thừa (vì đối với hầu hết các toán tử StableHLO, các loại đầu ra có thể được suy ra từ các đầu vào). Tuy nhiên, chữ ký op được cố ý đưa vào cú pháp StableHLO để tương thích với MLIR.

Dưới đây là một ví dụ về toán tử có ký hiệu là select_and_scatter. Hàm này sử dụng 3 giá trị đầu vào (%operand, %source và %init_value), 2 hàm đầu vào và 3 thuộc tính đầu vào (window_dimensions, window_strides và padding). Lưu ý cách chữ ký của toán tử chỉ bao gồm các loại giá trị đầu vào (nhưng không phải các loại hàm đầu vào và thuộc tính được cung cấp cùng dòng).

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

Hằng số

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

Hằng số ổn địnhHLO có một giá trị cố định và một kiểu cùng biểu thị một giá trị StableHLO. Nhìn chung, kiểu là một phần của cú pháp hằng số, trừ phi nó không rõ ràng (ví dụ: hằng số boolean rõ ràng có kiểu i1, trong khi một hằng số nguyên có thể có nhiều kiểu).

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

Hằng số boolean biểu thị các giá trị boolean true và false. Hằng số Boolean có kiểu i1.

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

Hằng số số nguyên biểu thị các giá trị số nguyên thông qua các chuỗi sử dụng ký hiệu thập phân hoặc thập lục phân. Các cơ số khác, ví dụ như nhị phân hoặc bát phân, không được hỗ trợ. Hằng số số nguyên có các quy tắc ràng buộc sau:

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

Hằng số dấu phẩy động biểu thị các giá trị dấu phẩy động thông qua các chuỗi sử dụng ký hiệu thập phân hoặc ký hiệu khoa học. Ngoài ra, bạn có thể sử dụng ký hiệu thập lục phân để chỉ định trực tiếp các bit cơ bản ở định dạng dấu phẩy động của loại tương ứng. Hằng số dấu phẩy động có các quy tắc ràng buộc sau:

• (C1) Nếu sử dụng ký hiệu không phải hệ thập lục phân, is_wellformed(float_literal, float_type).
• (C2) Nếu sử dụng ký hiệu thập lục phân, size(hexadecimal_digits) = num_bits(float_type) / 4.

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

Hằng số phức biểu thị các giá trị phức bằng cách sử dụng danh sách phần thực (đến trước) và phần ảo (đến sau). Ví dụ: (1.0, 0.0) : complex<f32> đại diện cho 1.0 + 0.0i và (0.0, 1.0) : complex<f32> đại diện cho 0.0 + 1.0i. Thứ tự lưu trữ các phần này trong bộ nhớ được xác định bằng cách triển khai. Hằng số phức có các quy tắc ràng buộc sau:

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

Hằng số tensor biểu thị các giá trị tensor bằng cách sử dụng các danh sách lồng nhau được chỉ định thông qua ký hiệu NumPy. Ví dụ: dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> đại diện cho một giá trị tensor với ánh xạ sau đây từ chỉ mục đến phần tử: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Việc triển khai sẽ xác định thứ tự lưu trữ các phần tử này trong bộ nhớ. Hằng số tensor có các quy tắc ràng buộc sau:

• (C1) has_syntax(tensor_literal, element_type(tensor_type)), trong đó:
  • 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)), trong đó:
  • 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:]).
  • nếu không, false.

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

Hằng số tensor lượng tử hoá biểu thị các giá trị tensor lượng tử hoá bằng cách sử dụng cùng một ký hiệu như hằng số tensor, với các phần tử được chỉ định là hằng số của kiểu lưu trữ. Hằng số tensor lượng tử hoá có các quy tắc ràng buộc sau:

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

Chuỗi cố định bao gồm các byte được chỉ định bằng ký tự ASCII và trình tự thoát. Các byte này không phụ thuộc vào phương thức mã hoá, vì vậy, việc diễn giải các byte này là do phương thức triển khai xác định. Hằng chuỗi có kiểu string.

Hoạt động

abs

Ngữ nghĩa

Thực hiện phép toán abs theo phần tử trên tensor operand và tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy thực hiện các thao tác sau:

• Đối với số nguyên có dấu: mô-đun số nguyên.
• Đối với số thực: abs từ IEEE-754.
• Đối với số phức: mô-đun số phức.
• Đối với các loại được lượng tử hoá: dequantize_op_quantize(abs, operand, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) shape(result) = shape(operand).
• (C2) baseline_element_type(result) được xác định là:
  • complex_element_type(element_type(operand)) nếu is_complex(operand).
  • baseline_element_type(operand).

Ví dụ

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

 Các ví dụ khác

thêm

Ngữ nghĩa

Thực hiện phép cộng theo phần tử của hai tensor lhs và rhs, đồng thời tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy làm như sau:

• Đối với boolean: logic OR.
• Đối với số nguyên: phép cộng số nguyên.
• Đối với số thực độ chính xác đơn: addition từ IEEE-754.
• Đối với số phức: phép cộng số phức.
• Đối với các kiểu lượng tử hoá: dequantize_op_quantize(add, lhs, rhs, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• Nếu toán tử sử dụng tensor không được lượng tử hoá:
  • (C1) type(lhs) = type(rhs) = type(result).
• Nếu phép toán sử dụng tensor lượng tử hoá:
  • (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) Nếu là is_per_axis_quantized(lhs), thì quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) = quantization_dimension(result).

Ví dụ

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

Ví dụ khác

after_all

Ngữ nghĩa

Đảm bảo rằng các thao tác tạo ra inputs được thực thi trước mọi thao tác phụ thuộc vào result. Việc thực thi thao tác này không làm gì cả, chỉ tồn tại để thiết lập các phần phụ thuộc dữ liệu từ result đến inputs.

Thông tin đầu vào

Kết quả đầu ra

Ví dụ

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

 Các ví dụ khác

all_gather

Ngữ nghĩa

Trong mỗi nhóm quy trình ở lưới quy trình StableHLO, hãy nối các giá trị của tensor operands từ mỗi quy trình dọc theo all_gather_dim và tạo ra tensor results.

Thao tác này chia lưới quy trình StableHLO thành process_groups được định nghĩa như sau:

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

Sau đó, trong mỗi process_group:

• operands...@receiver = [operand@sender for sender in process_group] cho tất cả receiver trong process_group.
• results...@process = concatenate(operands...@process, all_gather_dim) cho tất cả process trong process_group.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) 0 <= all_gather_dim < rank(operands...).
• (C2) is_unique(replica_groups).
• (C3) size(replica_groups) được xác định là:
  • num_replicas nếu bạn sử dụng cross_replica.
  • num_replicas nếu bạn sử dụng cross_replica_and_partition.
  • num_processes nếu bạn sử dụng flattened_ids.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Nếu use_global_device_ids = true, thì channel_id > 0.
• (C6) type(results...) = type(operands...) ngoại trừ:
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1).

Ví dụ

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

Ví dụ khác

all_reduce

Ngữ nghĩa

Trong mỗi nhóm quy trình trong lưới quy trình StableHLO, hãy áp dụng hàm rút gọn computation cho các giá trị của tensor operands từ mỗi quy trình và tạo ra tensor results.

Thao tác này chia lưới quy trình StableHLO thành process_groups được định nghĩa như sau:

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

Sau đó, trong mỗi process_group:

• results...@process[result_index] = exec(schedule) cho một số cây nhị phân schedule, trong đó:
  • exec(node) = computation(exec(node.left), exec(node.right)).
  • exec(leaf) = leaf.value.
• schedule là một cây nhị phân được xác định bằng cách triển khai, trong đó hoạt động duyệt theo thứ tự là to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0])).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) is_unique(replica_groups).
• (C2) size(replica_groups) được xác định là:
  • num_replicas nếu dùng cross_replica.
  • num_replicas nếu bạn sử dụng cross_replica_and_partition.
  • num_processes nếu bạn sử dụng flattened_ids.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Nếu use_global_device_ids = true, thì channel_id > 0.
• (C5) computation có loại (tensor<E>, tensor<E>) -> (tensor<E>) trong đó is_promotable(element_type(operand), E).
• (C6) shape(results...) = shape(operands...).
• (C7) element_type(results...) = E.

Ví dụ

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

 Các ví dụ khác

all_to_all

Ngữ nghĩa

Trong mỗi nhóm quy trình trong lưới quy trình StableHLO, hãy chia các giá trị của tensor operands dọc theo split_dimension thành các phần, phân tán các phần đã chia giữa các quy trình, nối các phần đã phân tán dọc theo concat_dimension và tạo tensor results. Thao tác này chia lưới quy trình StableHLO thành process_groups được định nghĩa như sau:

• cross_replica(replica_groups) nếu channel_id <= 0.
• cross_partition(replica_groups) nếu channel_id > 0.

Sau đó, trong mỗi process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) cho tất cả sender trong process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] trong đó receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (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) được xác định là:
  • num_replicas nếu dùng cross_replica.
  • num_partitions nếu dùng cross_partition.
• (C7) 0 <= replica_groups < size(replica_groups).
• (C8) dim(replica_groups, 1) = split_count.
• (C9) type(results...) = type(operands...) ngoại trừ, nếu split_dimension != concat_dimension:
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count.
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count.

Ví dụ

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

 Các ví dụ khác

và

Ngữ nghĩa

Thực hiện phép AND theo phần tử của hai tensor lhs và rhs, đồng thời tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy làm như sau:

• Đối với boolean: AND logic.
• Đối với số nguyên: thao tác bit AND.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

atan2

Ngữ nghĩa

Thực hiện phép toán atan2 theo phần tử trên tensor lhs và rhs, đồng thời tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy làm như sau:

• Đối với số thực: atan2 từ IEEE-754.
• Đối với số phức: atan2 phức.
• Đối với các loại được lượng tử hoá: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

batch_norm_grad

Ngữ nghĩa

Tính toán độ dốc của một số đầu vào của batch_norm_training truyền ngược từ grad_output và tạo ra các tensor grad_operand, grad_scale và grad_offset. Chính thức hơn, phép toán này có thể được biểu thị dưới dạng một phép phân ly cho các phép toán StableHLO hiện có bằng cách sử dụng cú pháp Python như sau:

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

Đối với các loại được lượng tử hoá, hãy thực hiện 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)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scale và grad_offset có cùng baseline_element_type.
• (C3) operand, grad_output và grad_operand có cùng hình dạng.
• (C4) scale, mean, variance, grad_scale và grad_offset có cùng hình dạng.
• (C5) size(scale) = dim(operand, feature_index).

Ví dụ

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

Ngữ nghĩa

Chuẩn hoá tensor operand trên tất cả các phương diện ngoại trừ phương diện feature_index và tạo ra một tensor result. Theo cách chính thức hơn, thao tác này có thể được biểu thị dưới dạng một quá trình phân ly cho các thao tác StableHLO hiện có bằng cách sử dụng cú pháp Python như sau:

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)

Đối với các loại được lượng tử hoá, hãy thực hiện 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)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance và result có cùng 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).

Ví dụ

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

Ngữ nghĩa

Tính toán giá trị trung bình và phương sai trên tất cả các phương diện, ngoại trừ phương diện feature_index, đồng thời chuẩn hoá tensor operand tạo ra tensor output, batch_mean và batch_var. Chính thức hơn, thao tác này có thể được biểu thị dưới dạng quá trình phân ly cho các thao tác ổn định hiện có bằng cú pháp Python như sau:

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

Đối với các loại được lượng tử hoá, hãy thực hiện 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)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var và output có cùng 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).

Ví dụ

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

Ngữ nghĩa

Thực hiện thao tác truyền bit trên tensor operand và tạo tensor result trong đó các bit của toàn bộ tensor operand được diễn giải lại bằng cách sử dụng loại tensor result.

Chính thức hơn, với E = element_type(operand), E' = element_type(result) và R = rank(operand):

• Nếu là num_bits(E') < num_bits(E), bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]).
• Nếu là num_bits(E') > num_bits(E), bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :]).
• Nếu num_bits(E') = num_bits(E), bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1]).

bits trả về giá trị đại diện trong bộ nhớ của một giá trị nhất định và hành vi của giá trị này được xác định bằng cách triển khai vì cách trình bày chính xác của tensor được xác định bằng cách triển khai và cách trình bày chính xác của các loại phần tử cũng được xác định bằng cách triển khai.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) Cho E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) và R = rank(operand):
  • Nếu num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Nếu giá trị là num_bits(E') < num_bits(E):
  • rank(result) = R + 1.
  • dim(result, i) = dim(operand, i) cho tất cả 0 <= i < R.
  • dim(result, R) * num_bits(E') = num_bits(E).
  • Nếu num_bits(E') > num_bits(E):
  • rank(result) = R - 1.
  • dim(result, i) = dim(operand, i) cho tất cả 0 <= i < R.
  • dim(operand, R - 1) * num_bits(E) = num_bits(E').
• (C2) Nếu là is_complex(operand) or is_complex(result), thì is_complex(operand) and is_complex(result).

Ví dụ

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

Ví dụ khác

broadcast_in_dim

Ngữ nghĩa

Mở rộng các phương diện và/hoặc thứ hạng của một tensor đầu vào bằng cách sao chép dữ liệu trong tensor operand và tạo ra một tensor result. Chính thức hơn, result[result_index] = operand[operand_index] trong đó đối với tất cả d trong axes(operand):

• operand_index[d] = 0 nếu dim(operand, d) = 1.
• operand_index[d] = result_index[broadcast_dimensions[d]].

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) element_type(result) được xác định bằng:
  • element_type(operand), nếu !is_per_axis_quantized(operand).
  • element_type(operand) ngoại trừ quantization_dimension(operand), scales(operand) và zero_points(operand) có thể khác với quantization_dimension(result), scales(result) và zero_points(result) tương ứng, nếu không.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Đối với tất cả d trong axes(operand):
  • dim(operand, d) = 1 hoặc
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Nếu is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Nếu dim(operand, quantization_dimension(operand)) = 1, thì scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result))).

Ví dụ

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

 Các ví dụ khác

ốp lưng

Ngữ nghĩa

Tạo ra kết quả từ việc thực thi chính xác một hàm từ branches tuỳ thuộc vào giá trị của index. Chính thức hơn là result = selected_branch(), trong đó:

• selected_branch = branches[index] nếu 0 <= index < size(branches).
• selected_branch = branches[-1] nếu không.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

cbrt

Ngữ nghĩa

Thực hiện phép tính căn bậc ba theo phần tử trên tensor operand và tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy làm như sau:

• Đối với số thực độ chính xác đơn: rootn(x, 3) từ IEEE-754.
• Đối với số phức: căn bậc ba phức.
• Đối với các loại được lượng tử hoá: dequantize_op_quantize(cbrt, operand, type(result))

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

ceil

Ngữ nghĩa

Thực hiện số nguyên trên của tensor operand theo từng phần tử và tạo ra tensor result. Triển khai phép toán roundToIntegralTowardPositive theo thông số kỹ thuật IEEE-754. Đối với các loại được lượng tử hoá, hãy thực hiện dequantize_op_quantize(ceil, operand, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

cholesky

Ngữ nghĩa

Tính toán phép phân ly Cholesky của một loạt ma trận.

Chính thức hơn, đối với tất cả i trong index_space(result), result[i0, ..., iR-3, :, :] là một phép phân ly Cholesky của a[i0, ..., iR-3, :, :], ở dạng ma trận tam giác dưới (nếu lower là true) hoặc tam giác trên (nếu lower là false). Các giá trị đầu ra trong tam giác đối diện, tức là tam giác trên nghiêm ngặt hoặc tam giác dưới nghiêm ngặt tương ứng, được xác định bằng cách triển khai.

Nếu tồn tại i mà ma trận đầu vào không phải là ma trận Hermitian dương xác định, thì hành vi này không được xác định.

Đối với các loại được lượng tử hoá, hãy thực hiện dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

kẹp

Ngữ nghĩa

Giới hạn mọi phần tử của tensor operand giữa giá trị tối thiểu và tối đa và tạo ra tensor result. Chính thức hơn là result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), trong đó min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Đối với các loại được lượng tử hoá, thực hiện dequantize_op_quantize(clamp, min, operand, max, type(result)).

Việc áp đặt thứ tự trên số phức có liên quan đến ngữ nghĩa đáng ngạc nhiên, vì vậy, trong tương lai, chúng tôi dự định sẽ xoá tính năng hỗ trợ số phức cho toán tử này (#560).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

collective_broadcast

Ngữ nghĩa

Trong mỗi nhóm quy trình trong lưới quy trình StableHLO, hãy gửi giá trị của tensor operand từ quy trình nguồn đến các quy trình mục tiêu và tạo một tensor result.

Thao tác này chia lưới quy trình StableHLO thành process_groups được định nghĩa như sau:

• cross_replica(replica_groups) nếu channel_id <= 0.
• cross_partition(replica_groups) nếu channel_id > 0.

Sau đó, result@process được xác định bằng:

• operand@process_groups[i, 0] nếu tồn tại một i sao cho quá trình này nằm trong process_groups[i].
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) nếu không.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) is_unique(replica_groups).
• (C2) 0 <= replica_groups < N trong đó N được xác định là:
  • num_replicas nếu bạn sử dụng cross_replica.
  • num_partitions nếu bạn sử dụng cross_partition.
• (C3) type(result) = type(operand).

Ví dụ

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

Ngữ nghĩa

Trong mỗi nhóm quy trình trong lưới quy trình StableHLO, hãy gửi giá trị của tensor operand từ quy trình nguồn đến quy trình mục tiêu và tạo một tensor result.

Thao tác này chia lưới quy trình StableHLO thành process_groups được định nghĩa như sau:

• cross_replica(source_target_pairs) nếu channel_id <= 0.
• cross_partition(source_target_pairs) nếu channel_id > 0.

Sau đó, result@process được xác định bằng:

• operand@process_groups[i, 0], nếu tồn tại một i sao cho process_groups[i, 1] = process.
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) nếu không.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (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, trong đó N được xác định là:
  • num_replicas nếu bạn sử dụng cross_replica.
  • num_partitions nếu dùng cross_partition.
• (C5) type(result) = type(operand).

Ví dụ

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

Ví dụ khác

so sánh

Ngữ nghĩa

Thực hiện so sánh theo phần tử của các tensor lhs và rhs theo comparison_direction và compare_type, đồng thời tạo ra một tensor result.

Các giá trị của comparison_direction và compare_type có ngữ nghĩa sau:

Đối với các loại phần tử boolean và số nguyên:

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

Đối với các loại phần tử dấu phẩy động có compare_type = FLOAT, toán tử này sẽ triển khai các phép toán IEEE-754 sau:

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

Đối với các loại phần tử dấu phẩy động có compare_type = TOTALORDER, phương thức này sử dụng tổ hợp toán tử totalOrder và compareQuietEqual trong IEEE-754.

Đối với các loại phần tử phức tạp, việc so sánh theo thứ tự bảng chữ cái của các cặp (real, imag) được thực hiện bằng cách sử dụng comparison_direction và compare_type được cung cấp. Việc áp đặt thứ tự trên các số phức liên quan đến ngữ nghĩa đáng ngạc nhiên, vì vậy, trong tương lai, chúng tôi dự định sẽ xoá tính năng hỗ trợ cho các số phức khi comparison_direction là GE, GT, LE hoặc LT (#560).

Đối với các loại được lượng tử hoá. thực hiện dequantize_compare(lhs, rhs, comparison_direction).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs).
• (C2) shape(lhs) = shape(rhs) = shape(result).
• (C3) compare_type được xác định là:
  • SIGNED nếu is_signed_integer(element_type(lhs)).
  • UNSIGNED nếu is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs)).
  • FLOAT hoặc TOTALORDER nếu is_float(element_type(lhs)).
  • FLOAT nếu is_complex(element_type(lhs)).

Ví dụ

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

 Các ví dụ khác

phức tạp

Ngữ nghĩa

Thực hiện chuyển đổi theo phần tử thành một giá trị phức từ một cặp giá trị thực và giá trị ảo, lhs và rhs, đồng thời tạo ra một tensor result.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) type(lhs) = type(rhs).
• (C2) shape(result) = shape(lhs).
• (C3) element_type(result) có loại complex<E> trong đó E = element_type(lhs).

Ví dụ

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

Ví dụ khác

tổng hợp

Ngữ nghĩa

Đóng gói một phép toán được tạo thành (thành phần) từ các phép toán StableHLO khác, lấy inputs và composite_attributes và tạo ra results. Ngữ nghĩa của hoạt động được triển khai bằng thuộc tính decomposition. Bạn có thể thay thế hoạt động composite bằng quá trình phân rã mà không cần thay đổi ngữ nghĩa của chương trình. Trong trường hợp việc nội tuyến quá trình phân ly không cung cấp cùng một ngữ nghĩa toán tử, hãy ưu tiên sử dụng custom_call.

Trường version (mặc định là 0) được dùng để biểu thị khi ngữ nghĩa của một tổ hợp thay đổi.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

nối

Ngữ nghĩa

Nối inputs dọc theo chiều dimension theo cùng thứ tự như các đối số đã cho và tạo ra một tensor result. Chính thức hơn, result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], trong đó:

1. id = d0 + ... + dk-1 + kd.
2. d bằng dimension và d0, ... là kích thước phương diện thứ d của inputs.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

Ví dụ khác

hằng số

Ngữ nghĩa

Tạo một tensor output từ một hằng số value.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

chuyển đổi

Ngữ nghĩa

Thực hiện chuyển đổi phần tử thông minh từ loại phần tử này sang loại phần tử khác trên tensor operand và tạo ra tensor result.

Đối với các lượt chuyển đổi boolean-to-any-supported-type (boolean sang bất kỳ loại nào được hỗ trợ), giá trị false được chuyển đổi thành 0 và giá trị true được chuyển đổi thành 1. Đối với lượt chuyển đổi any-supported-type-to-boolean, một giá trị bằng 0 sẽ được chuyển đổi thành false và các giá trị khác 0 sẽ được chuyển đổi thành true. Hãy xem phần bên dưới để biết cách hoạt động của tính năng này đối với các loại phức tạp.

Đối với các lượt chuyển đổi liên quan đến số nguyên sang số nguyên, số nguyên sang dấu phẩy động hoặc dấu phẩy động sang dấu phẩy động, nếu giá trị nguồn có thể được biểu thị chính xác trong loại đích, thì giá trị kết quả sẽ là giá trị biểu thị chính xác đó. Nếu không, hành vi này sẽ được xác định sau (#180).

Đối với các lượt chuyển đổi liên quan đến dấu phẩy động sang số nguyên, phần thập phân sẽ bị cắt bớt. Nếu không thể biểu thị giá trị bị cắt bớt trong loại đích, thì hành vi sẽ là TBD (#180).

Việc chuyển đổi liên quan đến số phức sang số phức tuân theo hành vi tương tự của các lượt chuyển đổi dấu phẩy động sang dấu phẩy động để chuyển đổi phần thực và phần ảo.

Đối với các phép chuyển đổi complex-to-any-other-type (phức tạp sang bất kỳ loại nào khác) và any-other-type-to-complex (bất kỳ loại nào khác sang phức tạp), giá trị ảo nguồn sẽ bị bỏ qua hoặc giá trị ảo đích sẽ được đặt thành 0. Quá trình chuyển đổi phần thực tuân theo các phép chuyển đổi dấu phẩy động.

Về nguyên tắc, toán tử này có thể biểu thị việc giải lượng tử hoá (chuyển đổi từ tensor lượng tử hoá sang tensor thông thường), lượng tử hoá (chuyển đổi từ tensor thông thường sang tensor lượng tử hoá) và lượng tử hoá lại (chuyển đổi giữa các tensor lượng tử hoá), nhưng hiện tại, chúng ta có các toán tử chuyên dụng cho việc đó – uniform_dequantize cho trường hợp sử dụng đầu tiên và uniform_quantize cho trường hợp sử dụng thứ hai và thứ ba. Trong tương lai, hai toán tử này có thể được hợp nhất vào convert (#1576).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

Ví dụ khác

tích chập

Ngữ nghĩa

Tính toán các sản phẩm chấm giữa các cửa sổ của lhs và lát cắt của rhs rồi tạo ra result. Sơ đồ sau đây cho thấy cách các phần tử trong result được tính toán từ lhs và rhs bằng một ví dụ cụ thể.

Một cách chính thức hơn, hãy cân nhắc việc điều chỉnh lại khung hình của các dữ liệu đầu vào sau đây theo lhs để có thể hiện các cửa sổ của 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).

Quá trình định khung lại này sử dụng các hàm trợ giúp sau:

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

Nếu feature_group_count = 1 và batch_group_count = 1, thì đối với tất cả output_spatial_index trong index_space(dim(result, output_spatial_dimensions...)), result[result_shape(:, output_spatial_index, :)] = dot_product trong đó:

• 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]). Tính năng này có vẻ không được sử dụng nên trong tương lai, chúng tôi sẽ dự định loại bỏ tính năng này (#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]).

Nếu 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).

Nếu 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).

Đối với các loại được lượng tử hoá, hãy thực hiện 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)).

Đối với các loại lượng tử hoá kết hợp, hãy thực hiện 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).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (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) Cho 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) Cho 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) Cho trước 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) được xác định là:
  • dim(lhs, input_batch_dimension) / batch_group_count nếu result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) nếu result_dim = output_feature_dimension.
  • num_windows nếu không, trong đó:
  • 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.
• Nếu toán tử sử dụng tensor không được lượng tử hoá:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Nếu thao tác sử dụng tensor lượng tử hoá:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Nếu là is_per_axis_quantized(result), thì quantization_dimension(result) = output_feature_dimension.
  • Nếu is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Nếu is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result).
  • Nếu giá trị là !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Ví dụ

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

Ví dụ khác

cosin

Ngữ nghĩa

Thực hiện phép toán cosin theo phần tử trên tensor operand và tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy thực hiện các thao tác sau:

• Đối với số thực: cos từ IEEE-754.
• Đối với số phức: cosin phức.
• Đối với các kiểu lượng tử hoá: dequantize_op_quantize(cosine, operand, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

Ví dụ khác

count_leading_zeros

Ngữ nghĩa

Thực hiện đếm số lượng bit số 0 ở đầu theo phần tử trong tensor operand và tạo ra một tensor result.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

custom_call

Ngữ nghĩa

Đóng gói một thao tác do triển khai xác định call_target_name, thao tác này sẽ lấy inputs và called_computations và tạo ra results. Bạn có thể dùng has_side_effect, backend_config và api_version để cung cấp thêm siêu dữ liệu do phương thức triển khai xác định.

Hiện tại, toán tử này chứa một tập hợp siêu dữ liệu khá lộn xộn, phản ánh sự phát triển hữu cơ của toán tử tương ứng trong trình biên dịch XLA. Trong tương lai, chúng tôi dự định hợp nhất siêu dữ liệu này (#741).

Thông tin đầu vào

Kết quả đầu ra

Ví dụ

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

chia

Ngữ nghĩa

Thực hiện phép chia từng phần tử của tensor lhs bị chia và số chia rhs, đồng thời tạo ra tensor result. Tuỳ thuộc vào loại phần tử, hãy làm như sau:

• Đối với số nguyên: phép chia số nguyên tạo ra thương đại số với mọi phần phân số bị loại bỏ.
• Đối với số thực: division từ IEEE-754.
• Đối với số phức: phép chia số phức.
• Đối với các loại được lượng tử hoá:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

Ví dụ khác

dot_general

Ngữ nghĩa

Tính toán tích vô hướng giữa các lát cắt của lhs và các lát cắt của rhs, đồng thời tạo ra một tensor result.

Chính thức hơn, result[result_index] = dot_product, trong đó:

• 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 trong đó size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) và 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)).

Đối với các loại được lượng tử hoá, hãy thực hiện 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)).

Đối với các kiểu lượng tử hoá kết hợp, thực hiện 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 kiểm soát sự đánh đổi giữa tốc độ và độ chính xác cho các phép tính trên phần phụ trợ của trình tăng tốc. Đây có thể là một trong những giá trị sau (hiện tại, ngữ nghĩa của các giá trị enum này chưa được chỉ định rõ, nhưng chúng tôi dự định sẽ giải quyết vấn đề này trong #755):

• DEFAULT: Tính toán nhanh nhất, nhưng có độ chính xác gần đúng nhất với số ban đầu.
• HIGH: Tính toán chậm hơn, nhưng gần đúng hơn với số ban đầu.
• HIGHEST: Phương thức tính toán chậm nhất, nhưng có độ chính xác gần nhất với số ban đầu.

DotAlgorithm xác định các thuộc tính chính của thuật toán dùng để triển khai phép tính dấu chấm, đồng thời xác định độ chính xác. Nếu bạn đặt các trường thuộc tính của thuật toán, thì precision_config phải là DEFAULT. DotAlgorithms không có giá trị mặc định, vì các tham số mặc định được xác định trong quá trình triển khai. Do đó, tất cả các trường thuật toán dấu chấm có thể được đặt thành None để chỉ định một thuật toán dấu chấm trống. Thay vào đó, thuật toán này sẽ sử dụng giá trị precision_config.

Các trường DotAlgorithm bao gồm:

• lhs_precision_type và rhs_precision_type, độ chính xác mà bên trái và bên phải của phép toán được làm tròn. Các loại độ chính xác độc lập với loại bộ nhớ của dữ liệu đầu vào và đầu ra.
• accumulation_type độ chính xác dùng để tích luỹ.
• lhs_component_count, rhs_component_count và num_primitive_operations áp dụng khi chúng ta thực hiện thuật toán phân tách LHS và/hoặc RHS thành nhiều thành phần và thực hiện nhiều thao tác dấu chấm "gốc" trên các giá trị đó – thường để mô phỏng độ chính xác cao hơn (ví dụ: Tận dụng Loại dữ liệu trí tuệ nhân tạo bfloat16 để tính toán độ chính xác cao hơn: bf16_6x tf). Đối với các thuật toán không có quá trình phân ly, bạn nên đặt các giá trị này thành 1.
• allow_imprecise_accumulation để chỉ định xem có được phép tích luỹ ở độ chính xác thấp hơn trong một số bước hay không (ví dụ: CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Ví dụ về thuộc tính 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}

Việc quyết định tổ hợp nào được hỗ trợ là tuỳ thuộc vào cách triển khai. Nhìn chung, không đảm bảo rằng mỗi thuật toán đều được người dùng StableHLO hỗ trợ trên mỗi loại trình tăng tốc. Nếu một thuật toán nhất định không được hỗ trợ, thì lỗi sẽ được báo cáo thay vì quay lại một phương án thay thế. Tính năng xác minh StableHLO sẽ cung cấp tính năng xác minh tốt nhất, ngăn chặn các thuật toán không được hỗ trợ trên bất kỳ phần cứng nào.

Hãy xem xla_data.proto > Algorithm để biết một số giá trị thuật toán được hỗ trợ. Phiếu yêu cầu hỗ trợ #2483 ghi lại kế hoạch tạo tài liệu tập trung về các thuật toán được hỗ trợ bởi phần phụ trợ.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (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).
• Nếu toán tử sử dụng tensor không được lượng tử hoá:
  • (C13) element_type(lhs) = element_type(rhs).
• Nếu phép toán sử dụng tensor lượng tử hoá:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) không có trong rhs_contracting_dimensions.
  • Nếu is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Nếu is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result).
  • Nếu !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Nếu giá trị là !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.

Ví dụ

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

 Các ví dụ khác

dynamic_broadcast_in_dim

Ngữ nghĩa

Toán tử này giống hệt về chức năng với toán tử broadcast_in_dim, nhưng hình dạng kết quả được chỉ định động thông qua output_dimensions.

Toán tử này cũng chấp nhận các thuộc tính không bắt buộc known_expanding_dimensions, known_nonexpanding_dimensions để thể hiện kiến thức tĩnh về hành vi mở rộng của các phương diện. Nếu bạn không chỉ định, tất cả phương diện đều được giả định là có thể mở rộng.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) element_type(result) được xác định bằng:
  • element_type(operand), nếu !is_per_axis_quantized(operand).
  • element_type(operand) ngoại trừ quantization_dimension(operand), scales(operand) và zero_points(operand) có thể khác với quantization_dimension(result), scales(result) và zero_points(result) tương ứng, nếu không.
• (C2) size(broadcast_dimensions) = rank(operand).
• (C3) 0 <= broadcast_dimensions < rank(result).
• (C4) is_unique(broadcast_dimensions).
• (C5) Đối với tất cả d trong axes(operand):
  • dim(operand, d) = 1 hoặc
  • dim(operand, d) = dim(result, broadcast_dimensions[d]).
• (C6) Nếu is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)].
  • Nếu dim(operand, quantization_dimension(operand)) = 1, thì 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).

Ví dụ

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

 Các ví dụ khác

dynamic_conv

Ngữ nghĩa

Toán tử này giống hệt về chức năng với toán tử convolution (tích chập), nhưng khoảng đệm được chỉ định động thông qua padding.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (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) Cho 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) Cho 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) Cho trước 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) được xác định là:
  • dim(lhs, input_batch_dimension) / batch_group_count nếu result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) nếu result_dim = output_feature_dimension.
  • num_windows nếu không, trong đó:
  • 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.
• Nếu toán tử sử dụng tensor không được lượng tử hoá:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Nếu thao tác sử dụng tensor lượng tử hoá:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Nếu là is_per_axis_quantized(result), thì quantization_dimension(result) = output_feature_dimension.
  • Nếu is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Nếu is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result).
  • Nếu giá trị là !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Ví dụ

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