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

StableHLO là một tập hợp thao tác cho các hoạt động cấp cao (HLO) trong máy mô hình học máy (ML). StableHLO hoạt động như một lớp di chuyển giữa các cấu hình Khung học máy và trình biên dịch học máy: các khung học máy tạo ra chương trình ổn định HLO tương thích với các trình biên dịch ML sử dụng các 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ạo ra Khả năng tương tác giữa nhiều khung máy học (chẳng hạn như TensorFlow, JAX và PyTorch) và trình biên dịch ML (chẳng hạn như XLA và IREE). Để đạt được mục tiêu này, cung cấp thông số kỹ thuật cho ngôn ngữ lập trình StableHLO.

Bản đặc tả này bao gồm ba phần chính. Đầu tiên, Phần Programs (Chương trình) mô tả cấu trúc của các chương trình StableHLO bao gồm các hàm StableHLO mà bản thân nó bao gồm các hoạt động ổn định HLO. Trong cấu trúc đó, phần Ops chỉ định ngữ nghĩa của hoạt động riêng lẻ. Phần Execution (Thực thi) cung cấp ngữ nghĩa cho tất cả các hoạt động 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ộ phần đặc điểm kỹ thuật.

Để xem thông số kỹ thuật của bản phát hành StableHLO trước, 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ụ: StableHLO v0.19.0 Spec. Để xem những thay đổi xảy ra ở mỗi lần tăng 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 ổn định HLO bao gồm một số lượng hàm ổn định HLO tuỳ ý. Dưới đây là một chương trình mẫu có hàm @main có 3 dữ liệu đầu vào (%image, %weights và %bias) và 1 đầu ra. Phần nội dung của hàm có 6 hoạt động.

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}

Các hàm ổn định (còn được gọi là hàm được đặt tên) có giá trị nhận dạng, đầu vào/đầu ra và nội dung. Trong tương lai, chúng tôi dự định giới thiệu siêu dữ liệu bổ sung 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 chương trình ngôn ngữ, với hai đặc điểm: 1) tất cả các giá trị nhận dạng đều có các dấu hiệu phân biệt các loại mã nhận dạng khác nhau, 2) mã nhận dạng giá trị hoàn toàn bằng 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

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

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

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

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

Trong tương lai, chúng tôi dự định khám phá việc mở rộng các loại tensor ngoài kích thước kích thước và loại phần tử, ví dụ: để bao gồm bố cục (#629) và độ thưa (#1078).

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

Các loại phần tử được định lượng biểu thị các giá trị số nguyên của một loại bộ nhớ trong phạm vi từ storage_min đến storage_max (bao gồm) tương ứng với giá trị dấu phẩy động thuộc một kiểu được biểu thị. Đối với một giá trị số nguyên đã cho 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 tham số lượng tử. storage_min và storage_max là không bắt buộc trong ngữ pháp, nhưng có các giá trị mặc định là min_value(storage_type) và max_value(storage_type). Các loại phần tử lượng tử hoá có những hạn chế 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à hằng số dấu phẩy động, nhưng có rất quan tâm đến thang đo số nguyên, được thể hiện bằng hệ số nhân và ca. Chúng tôi dự định sẽ khám phá tính năng này trong thời gian sắp tới (#1404).

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

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

Cuối cùng, chúng tôi dự định khám phá đại diện cho các thang đo chưa biết và số không tương tự như cách chúng ta dự định khám phá cách đại diện cho những giá trị chưa biết kích thước kích thước (#1407).

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

Trong tensor lượng tử hoá, lượng tử hoá có thể theo mỗi tensor, nghĩa là có một scale và zero_point cho toàn bộ tensor hoặc có thể là mỗi trục, tức là có nhiều scales và zero_points, một cặp trên mỗi lát cắt của một tham số cụ thể quantization_dimension. Trang trọng hơn khi dùng tensor t với lượng tử hoá theo mỗi trục, có dim(t, quantization_dimension) lát cắt trong quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :], Tất cả các phần tử trong lát i sử dụng scales[i] và zero_points[i] làm các tham số lượng tử của chúng. Sau đây là các loại tensor lượng tử hoá quy tắc ràng buộc:

• Đối với lượng tử mỗi tensor:
  • Không có hạn chế bổ sung.
• Đối với lượng tử trên 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ờ được tạo ra và tiêu thụ thực hiện một số phép toán. Mã thông báo được dùng để áp đặt thứ tự thực thi cho các thao tác như được mô tả trong phần Thực thi.

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

Loại bộ dữ liệu đại diện cho bộ dữ liệu, tức là danh sách không đồng nhất. Cặp đôi là di sản chỉ tồn tại để tương thích với HLO. Trong HLO, các bộ dữ liệu là được dùng để biểu thị nhiều giá trị đầu vào và đầu ra. Trong StableHLO, dữ liệu đầu vào đa dạng và đầu ra được hỗ trợ nguyên gốc và việc sử dụng duy nhất các bộ dữ liệu trong StableHLO là thể hiện toàn diện ABI HLO, ví dụ: T, tuple<T> và tuple<tuple<T>> có thể khác nhau đáng kể tuỳ thuộc vào một trong quá trình triển khai. Trong tương lai, chúng tôi dự định sẽ thay đổi ABI HLO tính năng này có thể cho phép chúng ta xoá các loại bộ dữ liệu khỏi 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 ::= 'f8E4M3FN' | 'f8E5M2' | 'f8E4M3FNUZ' | 'f8E5M2FNUZ'
            | 'f8E4M3B11FNUZ' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Loại phần tử đại diện cho phần tử thuộc loại tensor. Không giống như trong nhiều chương 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 ổn định HLO không thể biểu diễn trực tiếp giá trị của các loại này (do đó, thì gọi là biểu diễn các giá trị vô hướng thuộc loại T bằng tensor 0 chiều thuộc loại tensor<T>).

• Kiểu Boolean biểu thị các giá trị boolean true và false.
• Loại số nguyên có thể là đã ký (si) hoặc chưa ký (ui) và có một trong các độ rộng bit được hỗ trợ (2, 4, 8, 16, 32 hoặc 64). Loại siN đã ký biểu thị các giá trị số nguyên từ -2^(N-1) đến 2^(N-1)-1 các loại uiN không dấu biểu thị các giá trị số nguyên từ 0 đến 2^N-1.
• Loại dấu phẩy động có thể là một trong các loại sau:
  • các loại f8E4M3FN và f8E5M2 tương ứng với Mã hoá E4M3 và E5M2 của định dạng FP8 được mô tả trong Các định dạng của FP8 dành cho công nghệ học sâu.
  • Loại f8E4M3FNUZ và f8E5M2FNUZ tương ứng với E4M3 và E5M2 mã hoá các định dạng FP8 được mô tả trong Đị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 Chương trình đào tạo và suy luận dựa trên điểm nổi 8 bit kết hợp (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 được hiệu suất cao trên Cloud TPU.
  • Các loại f16, f32 và f64 tương ứng với binary16 ("độ chính xác bán chính xác"), binary32 ("độ chính xác đơn") và Định dạng binary64 ("độ chính xác kép") như mô tả trong tiêu chuẩn IEEE 754.
  • Loại tf32 tương ứng với định dạng TensorFloat32 và chỉ hỗ trợ hạn chế trong Ổn định HLO.
• Kiểu phức tạp biểu thị các giá trị phức có phần thực và một phần ảo thuộc cùng một loại phần tử. Khu phức hợp được hỗ trợ loại 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. Chúng có loại đầu vào (danh sách các loại ở bên trái của ->) và loại đầu ra (danh sách các loại ở bên phải ->). Trong nhiều chương trình ngôn ngữ, 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 chương trình ngôn ngữ, 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ử của chương trình.

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

Hoạt động ổn định HLO (còn được gọi là hoạt động) biểu thị một tập hợp khép kín các thao tác cấp cao trong các mô hình học máy. Như đã thảo luận ở trên, Cú pháp StableHLO chủ yếu lấy cảm hứng từ MLIR, không nhất thiết là giải pháp thay thế công thái học, nhưng được cho là phù hợp nhất cho mục tiêu của StableHLO là tăng khả năng tương tác giữa các khung ML và trình biên dịch ML.

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

Hoạt động ổn địnhHLO (còn được gọi là hoạt động) đều 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ự xác định duy nhất một trong các hoạt động được hỗ trợ. Xem 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 hoạt động sử dụng đầu vào và tạo ra đầu ra. Dữ liệu đầ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à đầu vào (cũng được cung cấp cố định). Loại dữ liệu đầu vào và đầu ra được tiêu thụ và sản xuất bởi một op phụ thuộc vào khả năng ghi nhớ của nó. Ví dụ: add op sử dụng 2 giá trị đầu vào và tạo ra 1 giá trị đầu ra. So với, Hoạt động 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}

Các hàm đầu vào (còn được gọi là hàm ẩn danh) rất tương tự như các 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) thẻ không khai báo loại đầu ra (loại đầu ra là 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 chưa được sử dụng (xem Unused chính thức ở trên), có khả năng tương thích với MLIR. Trong MLIR, có một khái niệm chung hơn là "khu vực" có thể có nhiều "chặn" gồm các 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 đối với phiên bản chính thức Unused để có thể phân biệt chúng 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 là chưa được sử dụng (nhưng vẫn cò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 những thuộc tính được hỗ trợ hằng số. Thẻ Google là phương thức chính để chỉ định siêu dữ liệu tĩnh cho chương trình phần tử. Ví dụ: hoạt động concatenate sử dụng thuộc tính dimension để chỉ định phương diện nối các giá trị đầu vào của phương diện đó. Tương tự, hoạt động 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 môi trường tự nhiên đôi khi có 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 hấp thụ các thuộc tính này vào hệ thống StableHLO hoặc cấm chúng xuất hiện trong chương trình StableHLO. Trong lúc đó, đây là danh sách thuộc tính:

• 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ý Op bao gồm các loại của tất cả giá trị đầu vào (danh sách các loại trên phần bên trái của ->) và loại của tất cả giá trị đầu ra (danh sách ở bên phải ->). Nói một cách chính xác, các loại dữ liệu đầu vào dư thừa và các loại đầu ra hầu như cũng luôn thừa (vì đối với hầu hết các hoạt động ổn định HLO, loại đầu ra có thể được suy ra từ dữ liệu đầu vào). Tuy nhiên, vận hành chữ ký là một phần có chủ ý của cú pháp StableHLO để tương thích với MLIR.

Dưới đây là một ví dụ về op có ghi nhớ là select_and_scatter. Chiến dịch này tiêu thụ 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 ý rằng chữ ký của op chỉ bao gồm các loại giá trị đầu vào của nó (nhưng không phải là loại hàm và thuộc tính đầu vào đượ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ị giá trị StableHLO. Thông thường, kiểu dữ liệu này là một phần của cú pháp hằng số, ngoại trừ khi không rõ ràng (ví dụ: hằng số boolean rõ ràng có loại i1, trong khi 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. Boolean (logic) hằng số 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ố biểu thị giá trị số nguyên thông qua các chuỗi sử dụng số thập phân hoặc ký hiệu thập lục phân. Các cơ sở khác, ví dụ: nhị phân hoặc bát phân, không được hỗ trợ. Hằng số nguyên có các hạn chế 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 khoa học. Ngoài ra, ký hiệu thập lục phân có thể là được dùng để 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 hạn chế sau:

• (C1) Nếu sử dụng ký hiệu không phải 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 tạp biểu thị các giá trị phức tạp bằng cách sử dụng danh sách phần thực (đến trước) và phần tưởng tượng (đế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ự trong đó sau đó các phần được lưu trữ trong bộ nhớ sẽ được triển khai xác định. Hằng số phức tạp có những hạn chế 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 qua Ký hiệu NumPy. Ví dụ: dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> biểu thị một giá trị tensor với ánh xạ sau từ chỉ mục đến các phần tử: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Sau đó, thứ tự lưu trữ các phần tử này trong bộ nhớ là do triển khai xác định. Hằng số Tensor có các hạn chế 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 thì 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 ký hiệu dưới dạng hằng số tensor, với các phần tử được chỉ định là hằng số loại bộ nhớ. Hằng số tensor lượng tử hoá có các hạn chế 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))

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

Hoạt động

bụng

Ngữ nghĩa

Thực hiện phép toán abs đa dạng trên tensor operand và tạo ra một result tensor. 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 độ chính xác đơn: abs từ IEEE-754.
• Đối với số phức: mô-đun phức.
• Đối với các kiểu 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 định nghĩa 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]

Ví dụ khác

thêm

Ngữ nghĩa

Thực hiện phép cộng theo từng phần tử của hai tensor lhs và rhs 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 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 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 phép toán sử dụng tensor không lượng tử hoá:
  • (C1) type(lhs) = type(rhs) = type(result).
• Nếu thao tác này 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 là 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 khi bất kỳ các toán tử phụ thuộc vào result. Việc thực thi thao tác này không ảnh hưởng gì đến nó 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

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ị trong số tensor operands từ mỗi quá trình dọc theo all_gather_dim và tạo ra Các tensor results.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups, tức là được xác định 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 định nghĩa là:
  • num_replicas nếu dùng cross_replica.
  • num_replicas nếu dùng cross_replica_and_partition.
  • num_processes nếu dùng flattened_ids.
• (C4) 0 <= replica_groups < size(replica_groups).
• (C5) Nếu là 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 ở lưới quy trình StableHLO, hãy áp dụng mức giảm hàm computation thành các giá trị của tensor operands trong mỗi tiến trình và tạo ra tensor results.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups, tức là được xác định 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à cây nhị phân do triển khai xác định có theo thứ tự truyền tải 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 định nghĩa là:
  • num_replicas nếu dùng cross_replica.
  • num_replicas nếu dùng cross_replica_and_partition.
  • num_processes nếu dùng flattened_ids.
• (C3) 0 <= replica_groups < size(replica_groups).
• (C4) Nếu là 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]

Ví dụ khác

all_to_all

Ngữ nghĩa

Trong mỗi nhóm quy trình ở lưới quy trình StableHLO, chia tách các giá trị của các tensor operands dọc theo split_dimension thành các phần, phân tán phần phân tách các phần giữa các quy trình, nối các phần nằm rải rác dọc theo concat_dimension và tạo ra tensor results. Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups, tức là được xác định 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 định nghĩa 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]]

Ví dụ khác

và

Ngữ nghĩa

Thực hiện cấu trúc phần tử AND của hai tensor lhs và rhs và tạo ra một result tensor. 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 boolean: logic AND.
• Đố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]]

Ví dụ khác

atan2

Ngữ nghĩa

Thực hiện toán tử atan2 thông minh trên tensor lhs và rhs 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 độ chính xác đơn: atan2 từ IEEE-754.
• Đối với số phức: phức atan2.
• Đối với các kiểu 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]

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 lan truyền ngược từ grad_output và tạo grad_operand, grad_scale và grad_offset tensor. Một cách 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 tách để các hoạt động ổn định hiện có bằ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 kiểu lượng tử hoá, 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 một 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 chiều, ngoại trừ chiều feature_index và tạo ra một tensor result. Chính thức hơn, đây là có thể được biểu thị dưới dạng quá trình phân ly đối với các toán tử ổn định hiện có 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 kiểu lượng tử hoá, 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ừ feature_index và chuẩn hoá tensor operand tạo ra output, batch_mean và batch_var tensor. Một cách chính thức hơn, thao tác này có thể được biểu thị dưới dạng phân rã thành các hoạt động ổ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 kiểu lượng tử hoá, 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 một 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 phép toán bitcast trên tensor operand và tạo ra 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 của tensor result.

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

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

bits trả về giá trị biểu thị trong bộ nhớ của một giá trị đã cho, kèm theo hành vi của giá trị đó được triển khai vì cách biểu diễn chính xác tensor là do triển khai xác định và cách thể hiện chính xác các loại phần tử là triển khai xác định.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) Cho trước 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 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(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 thứ nguyên và/hoặc thứ hạng của 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. Một cách chính thức hơn, result[result_index] = operand[operand_index] cho tất cả d ở 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 cung cấp bởi:
  • 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) phản hồi, 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 là 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]
//            ]
//          ]

Ví dụ khác

ốp lưng

Ngữ nghĩa

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

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

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]

Ví dụ khác

cbrt

Ngữ nghĩa

Thực hiện phép toán căn bậc ba toàn phần 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 độ 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 kiểu 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]

Ví dụ khác

ceil

Ngữ nghĩa

Thực hiện tensor kiểm soát phần tử của tensor operand và tạo ra tensor result. Triển khai thao tác roundToIntegralTowardPositive từ IEEE-754 đặc điểm kỹ thuật. Đối với các kiểu lượng tử hoá, 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]

Ví dụ khác

cholesky

Ngữ nghĩa

Tính toán quá trình phân rã Cholesky của một lô ma trận.

Chính thức hơn, cho tất cả i trong index_space(result), result[i0, ..., iR-3, :, :] là quá trình phân huỷ Cholesky của a[i0, ..., iR-3, :, :], ở dạng tam giác dưới (nếu lower là true) hoặc ma trận tam giác trên (nếu lower là false). Giá trị đầu ra trong tam giác đối diện, tức là tam giác nhọn ở trên hoặc tam giác dưới nghiêm ngặt tương ứng đều được xác định theo phương thức triển khai.

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

Đối với các kiểu lượng tử hoá, 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

Gắn mọi phần tử của tensor operand từ giá trị tối thiểu đến giá trị tối đa và tạo ra tensor result. Chính thức hơn, 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 lượng tử hoá, thực hiện dequantize_op_quantize(clamp, min, operand, max, type(result)).

Việc áp đặt một thứ tự cho số phức liên quan đến ngữ nghĩa đáng kinh ngạc, vì vậy, trong tương lai, chúng tôi sẽ có kế hoạch ngừng hỗ trợ số phức cho thao tác 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]

Ví dụ khác

collective_broadcast

Ngữ nghĩa

Trong mỗi nhóm quy trình ở 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 ra một Tensor result.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups, tức là được xác định như sau:

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

Sau đó, result@process được cung cấp bởi:

• operand@process_groups[i, 0] nếu tồn tại i để quá trình này 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 định nghĩa là:
  • num_replicas nếu dùng cross_replica.
  • num_partitions nếu 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 ở lưới quy trình StableHLO, hãy gửi giá trị của Tensor operand từ quá trình nguồn đến tiến trình mục tiêu và tạo ra một Tensor result.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups, tức là được xác định 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 cung cấp bởi:

• operand@process_groups[i, 0], nếu tồn tại 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 định nghĩa là:
  • num_replicas nếu 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 các phần tử của tensor lhs và rhs theo comparison_direction và compare_type rồi tạo tensor result.

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

Đố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, hoạt động sẽ triển khai các hoạt động IEEE-754 sau đây:

• 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, tuỳ chọn hoạt động sử dụng tổ hợp toán tử totalOrder và compareQuietEqual từ IEEE-754.

Đối với các loại phần tử phức tạp, phép so sánh từ điển của các cặp (real, imag) là thực hiện bằng comparison_direction và compare_type được cung cấp. Việc áp đặt một thứ tự cho số phức liên quan đến ngữ nghĩa đáng kinh ngạc, vì vậy, trong tương lai, chúng tôi sẽ có kế hoạch ngừng hỗ trợ số phức khi comparison_direction ở trạng thái GE, GT, LE hoặc LT (#560).

Đối với các loại 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 định nghĩa 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]

Ví dụ khác

phức tạp

Ngữ nghĩa

Thực hiện chuyển đổi từng phần tử thành một giá trị phức từ một cặp giá trị thực và các giá trị ảo, lhs và rhs, đồng thời tạo ra 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 thao tác được tạo thành (bao gồm) các thao tác ổn định khác, lấy inputs và composite_attributes và tạo ra results. Chiến lược phát hành đĩa đơn ngữ nghĩa của hoạt động được triển khai bằng thuộc tính decomposition. Chiến lược phát hành đĩa đơn Có thể thay thế hoạt động composite bằng quá trình phân rã mà không cần thay đổi chương trình ngữ nghĩa. Trong trường hợp cùng dòng trong quá trình phân ly không cung cấp cùng một ngữ nghĩa op, 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 một thuộc tính tổng hợp ngữ nghĩa.

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>

Ví dụ khác

nối

Ngữ nghĩa

Nối inputs dọc theo phương diện dimension theo cùng thứ tự như phương diện đã cho đối số và tạo ra một tensor result. Một cách 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 kích thước thứ d trong tổng số 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 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]]

Ví dụ khác

chuyển đổi

Ngữ nghĩa

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

Đối với lượt chuyển đổi boolean-to-any-supported-type, giá trị false là chuyển đổi thành 0 và giá trị true được chuyển đổi thành 1. Cho 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 được chuyển đổi thành true. Xem phần dưới đây để biết cách thực hiện phù hợp 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 đến số nguyên, số nguyên đến dấu phẩy động hoặc floating-point-to-floating-point, nếu giá trị nguồn có thể chính xác được biểu thị trong loại đích đến, thì giá trị kết quả là chính xác đại diện. Nếu không, hành vi sẽ được xác định sau (#180).

Đối với các lượt chuyển đổi liên quan đến floating-point-to-integer, phần phân số là bị cắt bớt. Nếu không thể biểu thị giá trị bị cắt bớt trong loại đích đến, hành vi sẽ được xác định sau (#180).

Việc chuyển đổi liên quan đến từ số phức sang số phức sẽ tuân theo hành vi tương tự của chuyển đổi floating-point-to-floating-point để chuyển đổi thực và phần ảo.

Đối với lượt chuyển đổi complex-to-any-other-type và complex-to-any-other-type, giá trị ảo nguồn bị bỏ qua hoặc giá trị ảo đích là có giá trị tương ứng là 0. Việc chuyển đổi phần thực tuân theo chuyển đổi dấu phẩy động.

Theo nguyên tắc, thao tác này có thể thể hiện việc loại bỏ lượng tử (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 tensor thành tensor lượng tử hoá) và lượng tử hoá (chuyển đổi giữa lượng tử hoá) tensor), nhưng hiện tại chúng tôi có hoạt động chuyên biệt 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 trường hợp sử dụng thứ hai và thứ ba. Trong tương lai, hai hoạt động 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 các tích vô hướng 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ơ đồ dưới đây cho thấy cách các phần tử trong result được tính toán từ lhs và rhs trong một ví dụ cụ thể.

Để cụ thể hơn, hãy cân nhắc việc điều chỉnh lại khung hình sau đây của các thông tin đầu vào về mặt lhs để có thể biểu thị 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).

Việc định lại khung hình này sử dụng các chức năng 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ì 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 dường như không được sử dụng, vì vậy trong tương lai, chúng tôi sẽ có kế hoạch loại bỏ nó (#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 giá trị là 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 giá trị là 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 kiểu lượng tử hoá, 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 kiểu lượng tử hoá kết hợp, 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 trước 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 trước 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 định nghĩa 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.
  • Nếu không là num_windows, 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 phép toán sử dụng tensor không lượng tử hoá:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Nếu thao tác này 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), sau đó là 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 giá trị là is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Nếu là 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 độ chính xác đơn: 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ố phần tử theo số phần tử số 0 ở đầu trong operand tensor và tạo ra 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]]

Ví dụ khác

custom_call

Ngữ nghĩa

Đóng gói thao tác triển khai call_target_name sẽ nhận inputs và called_computations, đồng thời tạo results. has_side_effect, backend_config và api_version có thể được dùng để cung cấp thêm siêu dữ liệu do triển khai xác định.

Hiện tại, hoạt động này chứa một tập hợp khá vô tổ chức siêu dữ liệu phản ánh sự tiến hoá tự nhiên của hoạt động tương ứng của nó 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 số bị chia lhs và tensor rhs của số chia và tạo ra 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: phép chia số nguyên tạo ra thương số đại số với bất kỳ đã loại bỏ phần thập phân.
• Đối với số thực độ chính xác đơn: division từ IEEE-754.
• Đối với số phức: phép chia phức.
• Đối với các loại 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 tích các chấm giữa các lát cắt lhs và lát cắt của rhs, từ đó tạo ra một Tensor result.

Chính thức hơn là 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 kiểu lượng tử hoá, 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ự cân bằng giữa tốc độ và độ chính xác cho các tính toán trên chương trình phụ trợ của trình tăng tốc. Đây có thể là một trong những cách sau (ở thời điểm, ngữ nghĩa của các giá trị enum này không được chỉ định rõ, nhưng chúng tôi đang dự định giải quyết vấn đề này trong #755):

• DEFAULT: Kết quả tính toán nhanh nhất, nhưng kém chính xác nhất so với kết quả ước tính số điện thoại ban đầu.
• HIGH: Phép tính toán chậm hơn, nhưng giá trị ước tính chính xác hơn số điện thoại ban đầu.
• HIGHEST: Phép tính toán chậm nhất, nhưng chính xác nhất so với số điện thoại 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 toán tử dấu chấm cũng xác định độ chính xác. Nếu thuộc tính thuật toán các trường được đặt, thì precision_config phải là DEFAULT. DotAlgorithms không có giá trị mặc định, vì các thông số mặc định đang được triển khai xác định. Như vậy, 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à LHS và RHS của thao tác được làm tròn đến. Loại độ chính xác độc lập với các loại bộ nhớ đầ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 sẽ á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 "gốc" trên những trang web đó giá trị – thường để mô phỏng độ chính xác cao hơn (ví dụ: Sử dụng kiểu dữ liệu trí tuệ nhân tạo bfloat16 để tính toán độ chính xác cao hơn: bf16_6x tf32_3x, v.v.). Đối với các thuật toán không phân ly, các giá trị này phải được đặt thành 1.
• allow_imprecise_accumulation để chỉ định liệu dữ liệu tích luỹ có độ chính xác thấp hơn hay không được phép thực hiện một số bước (ví dụ: CUBLASLT_MATMUL_DESC_FAST_ACCUM).

Ví dụ về các 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}

Điều này phụ thuộc vào việc triển khai để quyết định kết hợp nào được hỗ trợ. Trong nói chung, không thể đảm bảo rằng mỗi thuật toán được hỗ trợ trên mỗi loại trình tăng tốc theo người tiêu dùng StableHLO. Nếu một thuật toán cụ thể không được lỗi sẽ được đưa ra thay vì quay lại sử dụng thay thế. Quy trình xác minh trên StableHLO sẽ nỗ lực xác minh một cách hiệu quả 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.

Xem xla_data.proto > Algorithm đối với một số giá trị thuật toán được hỗ trợ. Phiếu yêu cầu hỗ trợ số 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ợ theo 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 phép toán sử dụng tensor không lượng tử hoá:
  • (C13) element_type(lhs) = element_type(rhs).
• Nếu thao tác này 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 là is_per_axis_quantized(rhs), thì quantization_dimension(rhs) không có trong rhs_contracting_dimensions.
  • Nếu giá trị là is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Nếu là is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result)
  • Nếu giá trị là !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]]
//          ]

Ví dụ khác

dynamic_broadcast_in_dim

Ngữ nghĩa

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

Thao tác này cũng chấp nhận các thuộc tính không bắt buộc known_expanding_dimensions, known_non_expanding_dimensions để thể hiện kiến thức tĩnh về hành vi mở rộng của thứ nguyên. Nếu bạn không chỉ định, tất cả cá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 cung cấp bởi:
  • 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) phản hồi, 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 là 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_non_expanding_dimensions).
• (C9) 0 <= known_expanding_dimensions < rank(operand).
• (C10) 0 <= known_non_expanding_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_non_expanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

Ví dụ khác

dynamic_conv

Ngữ nghĩa

Hoạt động này có chức năng giống với tích chập , nhưng khoảng đệm được chỉ định một cách linh độ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 trước 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 trước 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 định nghĩa 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.
  • Nếu không là num_windows, 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 phép toán sử dụng tensor không lượng tử hoá:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Nếu thao tác này 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), sau đó là 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 giá trị là is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Nếu là 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]]
//          ]]