Thông số kỹ thuật StableHLO

StableHLO là một bộ toán tử dành cho các hoạt động cấp cao (HLO) trong các mô hình học máy (ML). StableHLO hoạt động như một lớp hỗ trợ di chuyển giữa các khung máy học và trình biên dịch máy học: các khung máy học tạo ra các chương trình StableHLO tương thích với những trình biên dịch máy học 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ăng cường 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à 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.

Thông số kỹ thuật này bao gồm ba 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 StableHLO bao gồm các hàm StableHLO và bản thân các hàm đó cũng bao gồm các hoạt động StableHLO. Trong cấu trúc đó, phần Hoạt động chỉ định ngữ nghĩa của từng 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 cùng thực thi 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 dùng trong toàn bộ quy cách.

Chương trình

Program ::= {Func}

Các chương trình StableHLO bao gồm một số lượng hàm StableHLO tuỳ ý. Dưới đây là chương trình mẫu với hàm @main có 3 đầu vào (%image, %weights%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}

Hàm StableHLO (còn gọi là hàm được đặt tên) có 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 ra mắt 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 StableHLO 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ó dấu hiệu để phân biệt các loại giá trị nhận dạng, 2) giá trị nhận dạng có thể ở dạng số hoàn toàn để đơn giản hoá việc tạo các 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 loại giá trị (còn gọi là loại lớp đầu tiên) đại diện cho các giá trị StableHLO và loại không phải giá trị mô tả các phần tử khác của chương trình. 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 điểm đặc biệt chính là bản chất riêng theo miền của StableHLO, dẫn đến một số kết quả bất thường (ví dụ: 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 đại diện cho tensor, tức là mảng đa chiều. Chúng có hình dạngloại phần tử, trong đó hình dạng đại diện cho kích thước phương diện không âm theo thứ tự tăng dần của các 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à thứ 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. Biểu đồ này có hai chiều (hoặc nói cách khác là hai trục) – chiều thứ 0 và chiều thứ nhất – có kích thước là 2 và 3. Thứ hạng của mục này là 2.

Điều này xác định khả năng hỗ trợ cho hình dạng tĩnh, trong đó kích thước kích thước được xác định theo cách tĩnh. Trong tương lai, chúng tôi dự định cũng hỗ trợ hình dạng động, trong đó kích thước kích thước chưa được xác định một phần hoặc hoàn toàn (#8). Ngoài ra, chúng tôi đang có kế hoạch tìm hiểu cách 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
Tên Loại Giới hạn
storage_type loại số nguyên (C1–C4), (C9)
storage_min hằng số nguyên (C2), (C4), (C8)
storage_max hằng số nguyên (C3), (C4), (C8)
expressed_type loại dấu phẩy động (C1), (C5)
quantization_dimension hằng số nguyên không bắt buộc (C11–C13)
scales biến thiên số của hằng số dấu phẩy động (C5-C7), (C10), (C11), (C13)
zero_points biến thiên số của hằng số nguyên (C8–C10)

Các loại phần tử lượng tử 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 của một loại bộ nhớ được biểu thị. Đối với một giá trị số nguyên i nhất định, giá trị dấu phẩy động tương ứng f có thể được tính toán dưới dạng f = (i - zero_point) * scale, trong đó scalezero_point được gọi là tham số lượng tử. storage_minstorage_max là không bắt buộc trong ngữ pháp, nhưng có giá trị mặc định lần lượt là min_value(storage_type)max_value(storage_type). Các loại phần tử lượng tử có các quy tắc ràng buộc sau:

  • (C1) num_bits(storage_type) < num_bits(expressed_type).
  • (C2) type(storage_min) = storage_type.
  • (C3) type(storage_max) = storage_type.
  • (C4) min_value(storage_type) <= storage_min < storage_max <= max_value(storage_type).
  • (C5) type(scales...) = expressed_type.
  • (C6) 0 < scales.
  • (C7) is_finite(scales...).
  • (C8) storage_min <= zero_points <= storage_max.
  • (C9) type(zero_points...) = storage_type.
  • (C10) size(scales) = size(zero_points).
  • (C11) Nếu là is_empty(quantization_dimension), thì size(scales) = 1.
  • (C12) 0 <= quantization_dimension.

Hiện tại, QuantizationScale là một hằng số dấu phẩy động, nhưng người dùng rất quan tâm đến tỷ lệ dựa trên số nguyên, được biểu thị bằng hệ số và mức thay đổi. Chúng tôi dự định sẽ tìm hiểu về chủ đề này trong thời gian sắp tới (#1404).

Đ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 nhiều điểm 0 trong kiểu tensor 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 về 0 điểm 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 QuantizationStorageMinQuantizationStorageMax để xác định xem có nên áp dụng quy tắc ràng buộc nào cho các giá trị này và giá trị của tensor lượng tử hoá (#1406).

Cuối cùng, chúng tôi dự định tìm hiểu cách biểu thị các tỷ lệ và điểm 0 không xác định, tương tự như cách chúng tôi dự định khám phá cách biểu thị 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ị tensor với cá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ử này có các loại phần tử lượng tử hoá, thay vì các loại phần tử thông thường.

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

  • Đối với quá trình lượng tử hoá mỗi tensor:
    • Không có quy tắc ràng buộc bổ sung nào.
  • Đối với lượng tử hoá 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à sử dụng bởi một số thao tác. Mã thông báo được dùng để áp đặt lệnh thực thi đối với các thao tác như mô tả trong phần Thực thi.

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

Loại bộ đôi đại diện cho bộ dữ liệu, tức là danh sách không đồng nhất. Hàm 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, bộ dữ liệu được dùng để biểu thị các dữ liệu đầu vào và đầu ra biến thiên. Trong StableHLO, các dữ liệu đầu vào và đầu ra đa dạng được hỗ trợ sẵn, và việc sử dụng bộ dữ liệu duy nhất trong StableHLO là để biểu thị toàn diện ABI HLO, trong đó T, tuple<T>tuple<tuple<T>> có thể khác biệt đáng kể tuỳ thuộc vào một 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. Điều này có thể cho phép xoá các loại bộ dữ liệu khỏi StableHLO (#598).

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f8E4M3FN' | 'f8E5M2' | 'f8E4M3FNUZ' | 'f8E5M2FNUZ'
            | 'f8E4M3B11FNUZ' | 'bf16' | 'f16' | 'f32' | 'f64'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Loại phần tử biểu thị các phần tử thuộc loại tensor. Không giống như trong nhiều ngôn ngữ lập trình, các kiểu 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ị thuộc loại này (do đó, sẽ đặc trưng để biểu thị các giá trị vô hướng thuộc loại T với giá trị tensor 0 chiều thuộc loại tensor<T>).

  • Kiểu Boolean biểu thị các giá trị boolean truefalse.
  • Kiểu số nguyên có thể là số nguyên (si) hoặc không có chữ ký (ui) và có một trong các chiều rộng bit được hỗ trợ (4, 8, 16, 32 hoặc 64). Kiểu siN đã ký biểu thị giá trị số nguyên từ -2^(N-1) đến 2^(N-1)-1 bao gồm và loại uiN không có dấu biểu thị giá trị số nguyên từ 0 đến 2^N-1.
  • Các kiểu dấu phẩy động có thể là một trong những loại sau:
  • Loại phức tạp biểu thị các giá trị phức có một phần thực và một phần ảo thuộc cùng loại phần tử. Các kiểu 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 đối tượng này có loại dữ liệu đầu vào (danh sách các loại ở bên trái của ->) và loại dữ liệu đầu ra (danh sách các loại ở bên phải của ->). Trong nhiều ngôn ngữ lập trình, 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ị các chuỗi byte. Không giống như trong nhiều ngôn ngữ lập trình, kiểu 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 StableHLO (còn gọi là hoạt động) đại diện cho một tập hợp khép kín các hoạt động 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. Đây không nhất thiết phải là giải pháp thay thế hiệu quả nhất, nhưng có thể là lựa chọn phù hợp nhất với mục tiêu của StableHLO là tạo ra nhiều khả năng tương tác hơn giữa các khung máy học và trình biên dịch ML.

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

Các thao tác StableHLO (còn gọi là hoạt động) có tên, dữ liệu đầ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 những 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ợ.

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

  • builtin.module, func.func, func.callfunc.return (#425).
  • Toán tử chlo (#602).
  • Danh mục "Không nằm trong HLO" của các hoạt động StableHLO. ban đầu, chúng là một phần của tập hợp StableHLO nhưng sau đó bị coi là không phù hợp với nó: broadcast, create_token, cross-replica-sum, dot, einsum, torch_index_select, unary_einsum (#3).
  • Danh mục "Dynamism" của các hoạt động StableHLO. Các hoạt động này được tự khởi động từ MHLO, nhưng chúng tôi chưa chỉ định chúng: compute_reshape_shape, cstr_reshapable, dynamic_broadcast_in_dim, dynamic_conv, dynamic_gather, dynamic_iota, dynamic_pad, dynamic_reshape, real_dynamic_slice, set_dimension_size (#8).
  • Các phép tính hình dạng, bao gồm cả các toán tử arith, shapetensor (#8).
OpInputs        ::= OpInputValues OpInputFuncs OpInputAttrs
OpInputValues   ::= '(' [OpInputValue {',' OpInputValue}] ')'
OpInputValue    ::= ValueId
OpInputFuncs    ::= ['(' OpInputFunc {',' OpInputFunc} ')']
OpInputAttrs    ::= ['{' OpInputAttr {',' OpInputAttr} '}']
OpOutputs       ::= [OpOutput {',' OpOutput} '=']
OpOutput        ::= ValueId

Hoạt động vận hành sử dụng dữ liệu đầu vào và tạo đầ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à thuộc tính đầu vào (cũng được cung cấp theo cách tĩnh). Loại dữ liệu đầu vào và đầu ra mà một op sử dụng và tạo ra sẽ phụ thuộc vào khả năng ghi nhớ của op. Ví dụ: thao tác add sử dụng 2 giá trị đầu vào và tạo ra 1 giá trị đầu ra. Để so sánh, 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}

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

Cú pháp cho các hàm đầu vào có một phần hiện không được sử dụng (xem bản phát hành Unused ở trên) để tương thích với MLIR. Trong MLIR, có một khái niệm chung hơn về "khu vực" là "khu vực" có thể có nhiều "khối" hoạt động 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 được với nhau. StableHLO không có hoạt động nhảy nhanh, vì vậy, phần tương ứng của cú pháp MLIR không đượ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ị. Đây là một trong những hằng số được hỗ trợ. Đây là cách 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ụ: op concatenate sử dụng thuộc tính dimension để chỉ định kích thước mà các giá trị đầu vào được nối với nhau. Tương tự, hoạt động slice sử dụng nhiều thuộc tính như start_indiceslimit_indices để chỉ định các giới hạn dùng để phân tách giá trị đầu vào.

Hiện tại, các chương trình StableHLO ngoài tự nhiên đô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 sẽ đưa các thuộc tính này vào đối tượng kiểm thử StableHLO hoặc cấm chúng xuất hiện trong các chương trình StableHLO. Trong thời gian chờ đợi, sau đây là danh sách các thuộc tính sau:

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

Dưới đây là một ví dụ về cách ghi nhớ là select_and_scatter. Mô-đun này sử dụng 3 giá trị đầu vào (%operand, %source%init_value), 2 hàm đầu vào và 3 thuộc tính đầu vào (window_dimensions, window_stridespadding). Hãy lưu ý rằng chữ ký của op chỉ bao gồm các loại giá trị đầu vào (chứ không phải 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ố StableHLO có một kiểu cố định và một kiểu biểu thị giá trị StableHLO. Nhìn chung, kiểu dữ liệu là một phần của cú pháp hằng, trừ trường hợp rõ ràng (ví dụ: hằng số boolean không rõ ràng có kiểu 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 truefalse. Hằng số Boolean có loại 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ố nguyên biểu thị giá trị số nguyên thông qua những 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ị 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 những chuỗi sử dụng ký hiệu thập phân hoặc khoa học. Ngoài ra, bạn có thể 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 thập lục phân, hãy sử dụng is_wellformed(float_literal, float_type).
  • (C2) Nếu sử dụng ký hiệu thập lục phân, thì 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 danh sách một phần thực (đến trước) và một phần ảo (đến trước). Ví dụ: (1.0, 0.0) : complex<f32> đại diện cho 1.0 + 0.0i(0.0, 1.0) : complex<f32> đại diện cho 0.0 + 1.0i. Sau đó, thứ tự lưu trữ các phần này trong bộ nhớ sẽ do phương thức triển khai xác định. Hằng số phức tạp có các 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 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 có ánh xạ sau đây từ các 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ớ do triển khai xác đị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 ký hiệu như hằng số tensor, với các phần tử được chỉ định dưới dạng hằng số thuộc loại bộ nhớ của chúng. 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 cách sử dụng các ký tự ASCII và chuỗi 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 sẽ do quy trình 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 thao tác hấp thụ dựa trên 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ố 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 phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(abs, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của số nguyên có dấu, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1–C2)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại số nguyên hoặc dấu phẩy động đã ký hoặc tensor lượng tử hoá mỗi tensor (C1–C2)

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).
    • Nếu không, 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 việc cộng hai tensor lhsrhs theo cách tương tự, đồ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: toán tử logic OR.
  • Đối với số nguyên: cộng số nguyên.
  • Đối với số thực: addition từ IEEE-754.
  • Đối với số phức: phép cộng phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(add, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)
(I2) rhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Giới hạn

  • (C1) baseline_type(lhs) = baseline_type(rhs) = baseline_type(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 bất kỳ thao tác nào phụ thuộc vào result. Việc thực thi thao tác này không có tác dụng gì, 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

Nhãn Tên Loại
(I1) inputs số biến thiên của token

Kết quả đầu ra

Tên Loại
result token

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 trong lưới quy trình StableHLO, hãy nối các giá trị của tensor operand từ mỗi quy trình dọc theo all_gather_dim rồi tạo ra một tensor result.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups đượ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.
  • result@process = concatenate(operands@process, all_gather_dim) cho tất cả process trong process_group.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C6)
(I2) all_gather_dim hằng số có kiểu si64 (C1), (C6)
(I3) replica_groups Hằng số tensor 2 chiều thuộc loại si64 (C2–C4)
(I4) channel_id hằng số có kiểu si64 (C5)
(I5) use_global_device_ids hằng số có kiểu i1 (C5)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C6)

Giới hạn

  • (C1) 0 <= all_gather_dim < rank(operand).
  • (C2) is_unique(replica_groups).
  • (C3) size(replica_groups) được định nghĩa là:
    • num_replicas nếu sử dụng cross_replica.
    • num_replicas nếu sử dụng cross_replica_and_partition.
    • num_processes nếu sử 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(result) = type(operand) ngoại trừ:
    • dim(result, all_gather_dim) = dim(operand, all_gather_dim) * dim(process_groups, 1).

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2], [3, 4]]
// %operand@(1, 0): [[5, 6], [7, 8]]
%result = "stablehlo.all_gather"(%operand) {
  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<2x4xi64>
// %result@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]

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 một hàm rút gọn computation cho các giá trị của tensor operand từ mỗi quy trình 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 đượ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:

  • result@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 do triển khai xác định, có truyền tải theo thứ tự là to_destination_type(operands@process_group...[result_index], type(func_inputs(computation)[0])).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C5), (C6)
(I2) replica_groups số biến thiên của hằng số tensor 1 chiều thuộc kiểu si64 (C1–C3)
(I3) channel_id hằng số có kiểu si64 (C4)
(I4) use_global_device_ids hằng số có kiểu i1 (C4)
(I5) computation hàm (C5)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C6–C7)

Giới hạn

  • (C1) is_unique(replica_groups).
  • (C2) size(replica_groups) được định nghĩa là:
    • num_replicas nếu sử dụng cross_replica.
    • num_replicas nếu sử dụng cross_replica_and_partition.
    • num_processes nếu sử 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(result) = shape(operand).
  • (C7) element_type(result) = E.

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [1, 2, 3, 4]
// %operand@(1, 0): [5, 6, 7, 8]
%result = "stablehlo.all_reduce"(%operand) ({
  ^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_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<i64>) -> tensor<i64>
// %result@(0, 0): [6, 8, 10, 12]
// %result@(1, 0): [6, 8, 10, 12]

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 operand dọc theo split_dimension thành nhiều phần, phân tán các phần phân tách 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 ra một tensor result.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups đượ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(operand@sender, split_count, split_dimension) cho toàn bộ sender trong process_group.
  • scattered_parts@receiver = [split_parts@sender[receiver_index] for sender in process_group], trong đó receiver_index = process_group.index(receiver).
  • result@process = concatenate(scattered_parts@process, concat_dimension).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C3), (C9)
(I2) split_dimension hằng số có kiểu si64 (C1), (C2), (C9)
(I3) concat_dimension hằng số có kiểu si64 (C3), (C9)
(I4) split_count hằng số có kiểu si64 (C2), (C4), (C8), (C9)
(I5) replica_groups Hằng số tensor 2 chiều thuộc loại si64 (C5–C8)
(I6) channel_id hằng số có kiểu si64

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C9)

Giới hạn

  • (C1) 0 <= split_dimension < rank(operand).
  • (C2) dim(operand, split_dimension) % split_count = 0.
  • (C3) 0 <= concat_dimension < rank(operand).
  • (C4) 0 < split_count.
  • (C5) is_unique(replica_groups).
  • (C6) size(replica_groups) được định nghĩa là:
    • num_replicas nếu sử dụng cross_replica.
    • num_partitions nếu sử dụng cross_partition.
  • (C7) 0 <= replica_groups < size(replica_groups).
  • (C8) dim(replica_groups, 1) = split_count.
  • (C9) type(result) = type(operand) ngoại trừ:
    • dim(result, split_dimension) = dim(operand, split_dimension) / split_count.
    • dim(result, concat_dimension) = dim(operand, concat_dimension) * split_count.

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.all_to_all"(%operand) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
} : (tensor<2x4xi64>) -> tensor<4x2xi64>
// %result@(0, 0): [[1, 2],
//                  [5, 6],
//                  [9, 10],
//                  [13, 14]]
// %result@(1, 0): [[3, 4],
//                  [7, 8],
//                  [11, 12],
//                  [15, 16]]

Ví dụ khác

Ngữ nghĩa

Thực hiện phần tử AND của hai tensor lhsrhs, đồ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 AND.
  • Đối với số nguyên: bitwise AND.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của kiểu boolean hoặc số nguyên (C1)
(I2) rhs tensor của kiểu boolean hoặc số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu boolean hoặc số nguyên (C1)

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

atan2

Ngữ nghĩa

Thực hiện thao tác atan2 tương ứng trên tensor lhsrhs, đồ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ố thực: atan2 từ IEEE-754.
  • Đối với số phức: phức atan2.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(atan2, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) rhs tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

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 truyền ngược từ grad_output và tạo ra tensor grad_operand, grad_scalegrad_offset. Chính thức hơn, thao tác này có thể được biểu thị dưới dạng phân ly cho các thao tác StableHLO 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 loại lượng tử hoá, 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

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1–C3), (C5)
(I2) scale Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C4), (C5)
(I3) mean Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C4)
(I4) variance Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C4)
(I5) grad_output tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C2), (C3)
(I6) epsilon hằng số có kiểu f32
(I7) feature_index hằng số có kiểu si64 (C1), (C5)

Kết quả đầu ra

Tên Loại Giới hạn
grad_operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C2), (C3)
grad_scale Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C4)
grad_offset Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C4)

Giới hạn

  • (C1) 0 <= feature_index < rank(operand).
  • (C2) operand, scale, mean, variance, grad_output, grad_operand, grad_scalegrad_offset có cùng baseline_element_type.
  • (C3) operand, grad_outputgrad_operand có cùng hình dạng.
  • (C4) scale, mean, variance, grad_scalegrad_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 chiều, ngoại trừ chiều feature_index và tạo một tensor result. Chính thức hơn, thao tác này có thể được biểu thị dưới dạng phân ly các hoạt động 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 lượng tử hoá, 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

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1–C7)
(I2) scale Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C3)
(I3) offset Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C4)
(I4) mean Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C5)
(I5) variance Tensor 1 chiều thuộc loại lượng tử hoá dấu phẩy động hoặc mỗi tensor (C2), (C6)
(I6) epsilon hằng số có kiểu f32
(I7) feature_index hằng số có kiểu si64 (C1), (C3-C6)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C2), (C7)

Giới hạn

  • (C1) 0 <= feature_index < rank(operand).
  • (C2) operand, scale, offset, mean, varianceresult 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 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_meanbatch_var. Chính thức hơn, thao tác này có thể được biểu thị dưới dạng phân ly cho các hoạt động StableHLO hiện có bằng cách sử dụ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 lượng tử hoá, 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

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) scale Tensor 1 chiều của dấu phẩy động hoặc mỗi tensor lượng tử hoá (C2), (C3)
(I3) offset Tensor 1 chiều của dấu phẩy động hoặc mỗi tensor lượng tử hoá (C2), (C4)
(I4) epsilon hằng số có kiểu f32 (C1), (C3-C6)
(I5) feature_index hằng số có kiểu si64 (C1), (C3-C6)

Kết quả đầu ra

Tên Loại Giới hạn
output tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C7)
batch_mean Tensor 1 chiều của dấu phẩy động hoặc mỗi tensor lượng tử hoá (C2), (C5)
batch_var Tensor 1 chiều của dấu phẩy động hoặc mỗi tensor lượng tử hoá (C2), (C6)

Giới hạn

  • (C1) 0 <= feature_index < rank(operand).
  • (C2) operand, scale, offset, batch_mean, batch_varoutput 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 bitcast trên tensor operand và tạo ra một tensor result trong đó các bit của toàn bộ tensor operand sẽ được diễn giải lại bằng cách sử dụng loại tensor result.

Chính thức hơn, khi có E = element_type(operand), E' = element_type(result)R = rank(operand):

  • Nếu num_bits(E') < num_bits(E), bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1]).
  • Nếu 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ề nội dung biểu diễn trong bộ nhớ của một giá trị nhất định và hành vi của nó được xác định theo phương thức triển khai vì cách trình bày chính xác của tensor được xác định theo phương thức triển khai, và cách trình bày chính xác các loại phần tử cũng được xác định theo phương thức triển khai.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor hoặc tensor lượng tử hoá (C1–C2)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor hoặc tensor lượng tử hoá (C1–C2)

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)R = rank(operand):
    • Nếu num_bits(E') = num_bits(E), hãy shape(result) = shape(operand).
    • 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(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 kích thước 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. Chính thức hơn, result[result_index] = operand[operand_index], trong đó cho tất cả d trong axes(operand):

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

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor hoặc tensor lượng tử hoá (C1–C2), (C5–C6)
(I2) broadcast_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C2–C6)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor hoặc tensor lượng tử hoá (C1), (C3), (C5-C6)

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)zero_points(operand) có thể khác với quantization_dimension(result), scales(result)zero_points(result) tương ứ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 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, result = selected_branch(), trong đó:

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

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) index Tensor 0 chiều thuộc loại si32
(I2) branches số biến thiên của hàm (C1–C4)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên, tensor lượng tử hoá hoặc mã thông báo (C4)

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 toán tử căn bậc ba trên tensor operand và 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ố thực: 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 lượng tử hoá: dequantize_op_quantize(cbrt, operand, type(result))

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

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

câu lạc bộ ceil

Ngữ nghĩa

Thực hiện ceil theo phần tử của tensor operand và tạo ra một tensor result. Triển khai thao tác roundToIntegralTowardPositive từ quy cách IEEE-754. Đối với các loại lượng tử hoá, thực hiện dequantize_op_quantize(ceil, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

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 ly Cholesky của một loạt ma trận.

Nói một cách chính thức hơn, đối với mọi i trong index_space(result), result[i0, ..., iR-3, :, :] là một quá trình phân ly Cholesky của a[i0, ..., iR-3, :, :], ở dạng ma trận tam giác dưới (nếu lowertrue) hoặc ma trận tam giác trên (nếu lowerfalse). 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 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à ma trận xác định dương Hermitian, thì hành vi này không xác định.

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

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) a tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1–C3)
(I2) lower Hằng số tensor 0 chiều thuộc loại i1

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

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

Kẹp 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 một 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 lượng tử hoá, hãy thực hiện dequantize_op_quantize(clamp, min, operand, max, type(result)).

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

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) min Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C3)
(I2) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C4)
(I3) max Tensor lượng tử hoá tensor hoặc mỗi tensor (C2), (C3)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C4)

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 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 sẽ chia lưới quy trình StableHLO thành process_groups đượ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 sao cho quy trình này nằm trong process_groups[i].
  • Nếu không, broadcast_in_dim(constant(0, element_type(result)), [], type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor (C3)
(I2) replica_groups số biến thiên của hằng số tensor 1 chiều thuộc kiểu si64 (C1), (C2)
(I3) channel_id hằng số có kiểu si64

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor (C3)

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 sử dụng cross_replica.
    • num_partitions nếu 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 ra một tensor result.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups đượ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 một i sao cho process_groups[i, 1] = process.
  • Nếu không, broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C5)
(I2) source_target_pairs Hằng số tensor 2 chiều thuộc loại si64 (C1–C4)
(I3) channel_id hằng số có kiểu si64

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

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 sử dụng cross_replica.
    • num_partitions nếu sử 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 phần tử của tensor lhsrhs theo comparison_directioncompare_type, đồng thời tạo ra một tensor result.

Giá trị của comparison_directioncompare_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, hoạt động này sẽ triển khai các thao tác 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, hoạt động này sử dụng kết hợp các thao tác totalOrdercompareQuietEqual từ IEEE-754. Tính năng này có vẻ không được sử dụng, vì vậy, chúng tôi dự định xoá tính năng này trong tương lai (#584).

Đối với các loại phần tử phức tạp, việc so sánh từ vựng của các cặp (real, imag) được thực hiện bằng cách sử dụng comparison_directioncompare_type đã cung cấp. Việc áp đặt thứ tự trên các số phức liên quan đến ngữ nghĩa bất ngờ, vì vậy, trong tương lai, chúng tôi dự định ngừng hỗ trợ số phức khi comparison_directionGE, 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

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C3)
(I2) rhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C2)
(I3) comparison_direction enum của EQ, NE, GE, GT, LELT
(I4) compare_type enum của FLOAT, TOTALORDER, SIGNEDUNSIGNED (C3)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại boolean (C2)

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 là 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 việc chuyển đổi theo phần tử thành một giá trị phức tạp từ một cặp giá trị thực và ảo, lhsrhs, đồng thời tạo ra tensor result.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor thuộc loại f32 hoặc f64 (C1–C3)
(I2) rhs tensor thuộc loại f32 hoặc f64 (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor thuộc loại phức (C2), (C3)

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

concatenate

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 là result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1], trong đó:

  1. id = d0 + ... + dk-1 + kd.
  2. d bằng dimensiond0, ... là kích thước kích thước thứ d của inputs.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1–C6)
(I2) dimension hằng số có kiểu si64 (C2), (C4), (C6)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C5–C6)

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 value không đổi.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) value hằng số (C1)

Kết quả đầu ra

Tên Loại Giới hạn
output Tensor hoặc tensor lượng tử hoá (C1)

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 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 tensor result.

Đối với lượt chuyển đổi boolean-to-any-supported-type, giá trị false sẽ đượ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, hệ thống sẽ chuyển đổi giá trị 0 thành false và giá trị khác 0 sẽ được chuyển đổi thành true. Vui lòng xem phần bên dưới để biết cách công cụ này hoạt động đối với các kiểu phức tạp.

Đối với các lượt chuyển đổi liên quan đến số nguyên thành số nguyên, số nguyên thành số dấu phẩy động hoặc dấu phẩy động đến 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 đến, thì giá trị kết quả sẽ là giá trị biểu diễn chính xác đó. Nếu không, hành vi sẽ là TBD (#180).

Đối với các lượt chuyển đổi liên quan đến floating-point-to-integer, phần phân số 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 đến, thì hành vi sẽ là TBD (#180).

Phép chuyển đổi phức tạp thành phức tạp tuân theo hành vi tương tự của lượt chuyển đổi floating-point-to-floating-point cũng như để chuyển đổi các phần thực và ảo.

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

Về nguyên tắc, thao tác này có thể biểu thị quá trình 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 sang tensor lượng tử hoá) và tái lượng tử (chuyển đổi giữa tensor lượng tử hoá). Tuy nhiên, hiện tại chúng tôi đã thực hiện các thao tác dành riêng cho việc này – 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 hoạt động này có thể được hợp nhất vào convert (#1576).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor (C1)

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 sản phẩm chấm giữa các cửa sổ của lhs và các lát cắt của rhs, đồng thời tạo ra result. Sơ đồ sau đây cho thấy cách tính toán các phần tử trong result từ lhsrhs, sử dụng một ví dụ cụ thể.

Chính thức hơn, hãy cân nhắc việc thiết lập lại khung hình đầu vào sau đây theo 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 sửa lại khung hình này sử dụng các chức năng trợ giúp sau đây:

  • 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 = 1batch_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, vì vậy trong tương lai, chúng tôi dự định xoá (#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 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)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C10-C11), (C14) (C25), (C27-C30)
(I2) rhs Tensor hoặc tensor lượng tử hoá (C1), (C14-C16), (C25), (C27-C32)
(I3) window_strides Hằng số tensor 1 chiều thuộc loại si64 (C2–C3), (C25)
(I4) padding Hằng số tensor 2 chiều thuộc loại si64 (C4), (C25)
(I5) lhs_dilation Hằng số tensor 1 chiều thuộc loại si64 (C5-C6), (C25)
(I6) rhs_dilation Hằng số tensor 1 chiều thuộc loại si64 (C7-C8), (C25)
(I7) window_reversal Hằng số tensor 1 chiều thuộc loại i1 (C9)
(I8) input_batch_dimension hằng số có kiểu si64 (C10), (C13), (C25)
(I9) input_feature_dimension hằng số có kiểu si64 (C11), (C13-C14)
(I10) input_spatial_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C12), (C13), (C25)
(I11) kernel_input_feature_dimension hằng số có kiểu si64 (C14), (C18)
(I12) kernel_output_feature_dimension hằng số có kiểu si64 (C15-C16), (C18), (C25), (C32)
(I13) kernel_spatial_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C17-C18), (C25)
(I14) output_batch_dimension hằng số có kiểu si64 (C20), (C25)
(I15) output_feature_dimension hằng số có kiểu si64 (C20), (C25), (C33)
(I16) output_spatial_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C19-C20), (C25)
(I17) feature_group_count hằng số có kiểu si64 (C11), (C14), (C16), (C21), (C23)
(I18) batch_group_count hằng số có kiểu si64 (C10), (C15), (C22), (C23), (C25)
(I19) precision_config số enum biến đổi của DEFAULT, HIGHHIGHEST (C24)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor hoặc tensor lượng tử hoá (C25-C28), (C30-C31), (C33)

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) Đã cung cấp 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) Đã cung cấp 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) Đã cung cấp 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, hãy chọn 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 thao tác này 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_tensor(lhs) and is_quantized_tensor(rhs) and is_quantized_tensor(result).
    • (C29) storage_type(lhs) = storage_type(rhs).
    • (C30) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
    • (C31) Nếu là is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result).
    • (C32) Nếu là is_per_axis_quantized(rhs), thì quantization_dimension(rhs) = kernel_output_feature_dimension.
    • (C33) Nếu là is_per_axis_quantized(result), thì quantization_dimension(result) = output_feature_dimension.

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 = dense<4> : tensor<2xi64>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = dense<2> : tensor<2xi64>,
  rhs_dilation = dense<1> : tensor<2xi64>,
  window_reversal = dense<false> : tensor<2xi1>,
  // 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]>,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi32>, tensor<3x3x1x1xi32>) -> tensor<1x2x2x1xi32>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

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 làm như 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 loại lượng tử hoá: dequantize_op_quantize(cosine, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

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 việc đếm phần tử theo số lượng bit 0 đứng đầu trong tensor operand và tạo ra một tensor result.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của kiểu số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên (C1)

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 một thao tác được triển khai xác định call_target_name. Thao tác này sẽ lấy inputscalled_computations và tạo ra results. Có thể sử dụng has_side_effect, backend_configapi_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, thao tác này chứa một tập hợp siêu dữ liệu khá thiếu tổ chức, phản ánh sự phát triển tự nhiên của hoạt động đối tác 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

Nhãn Tên Loại
(I1) inputs số lượng giá trị biến thiên
(I2) call_target_name hằng số có kiểu string
(I3) has_side_effect hằng số có kiểu i1
(I4) backend_config hằng số có kiểu string
(I5) api_version hằng số có kiểu si32
(I6) called_computations biến thiên số của hằng số kiểu string

Kết quả đầu ra

Tên Loại
results số lượng giá trị biến thiên

Ví dụ

%results = "stablehlo.custom_call"(%input0) {
  call_target_name = "foo",
  has_side_effect = false,
  backend_config = "bar",
  api_version = 1 : i32,
  called_computations = [@foo]
} : (tensor<f64>) -> tensor<f64>

chia

Ngữ nghĩa

Thực hiện việc phân chia các phần tử tensor bị chia lhs và số chia 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ố nguyên: phép chia số nguyên tạo ra thương số đại số với phần phân số bất kỳ bị loại bỏ.
  • Đối với số thực: 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

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của số nguyên, dấu phẩy động hoặc loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) rhs tensor của số nguyên, dấu phẩy động hoặc loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

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 sản phẩm chấm 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 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)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 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)).

Thao tác này chỉ chỉ định ngữ nghĩa cho lượng tử hoá mỗi tensor. Quá trình lượng tử hoá trên mỗi trục đang diễn ra (#1574). Ngoài ra, trong tương lai, chúng tôi có thể cân nhắc việc thêm tính năng hỗ trợ lượng tử hoá kết hợp (#1575).

precision_config kiểm soát sự đánh đổi giữa tốc độ và độ chính xác trong 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 trường hợp 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 đang có kế hoạch giải quyết vấn đề này trong #755):

  • DEFAULT: Cách tính nhanh nhất nhưng cách tính gần đúng nhất nhưng ít chính xác nhất so với số gốc.
  • HIGH: Phương pháp tính toán chậm hơn nhưng cách tính xấp xỉ chính xác hơn với số ban đầu.
  • HIGHEST: Cách tính toán chậm nhất nhưng là cách tính gần đúng chính xác nhất với số ban đầu.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C5-C6), (C9-C10), (C12-C16)
(I2) rhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C7-C10), (C12)
(I3) lhs_batching_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C1), (C3), (C5), (C9), (C12)
(I4) rhs_batching_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C1), (C4), (C7), (C9)
(I5) lhs_contracting_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C2), (C3), (C6), (C10)
(I6) rhs_contracting_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C2), (C4), (C8), (C10)
(I7) precision_config số enum biến đổi của DEFAULT, HIGHHIGHEST (C11)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C12), (C14), (C16)

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 thao tác này 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) and is_quantized(rhs) and is_quantized(result).
    • (C15) storage_type(lhs) = storage_type(rhs).
    • (C16) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
    • (C17) zero_points(rhs) = 0.

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>]
} : (tensor<2x2x2xi64>, tensor<2x2x2xi64>) -> tensor<2x2x2xi64>
// %result: [
//           [[1, 2],
//            [3, 4]],
//           [[5, 6],
//            [7, 8]]
//          ]

Ví dụ khác

dynamic_slice

Ngữ nghĩa

Trích xuất một lát cắt từ operand bằng cách sử dụng các chỉ mục khởi động được tính toán động và tạo ra một tensor result. start_indices chứa các chỉ mục bắt đầu của lát cắt cho từng phương diện có thể được điều chỉnh và slice_sizes chứa kích thước của lát cắt cho mỗi phương diện. Chính thức hơn, result[result_index] = operand[operand_index], trong đó:

  • adjusted_start_indices = clamp(0, start_indices, shape(operand) - slice_sizes).
  • operand_index = adjusted_start_indices + result_index.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C2), (C4)
(I2) start_indices số biến thiên của tensor 0 chiều của kiểu số nguyên (C2), (C3)
(I3) slice_sizes Hằng số tensor 1 chiều thuộc loại si64 (C2), (C4), (C5)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C5)

Giới hạn

  • (C1) element_type(operand) = element_type(result).
  • (C2) size(start_indices) = size(slice_sizes) = rank(operand).
  • (C3) same(type(start_indices...)).
  • (C4) 0 <= slice_sizes <= shape(operand).
  • (C5) shape(result) = slice_sizes.

Ví dụ

// %operand: [
//            [0, 0, 1, 1],
//            [0, 0, 1, 1],
//            [0, 0, 0, 0],
//            [0, 0, 0, 0]
//           ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_slice"(%operand, %start_indices0, %start_indices1) {
  slice_sizes = dense<[2, 2]> : tensor<2xi64>
} : (tensor<4x4xi32>, tensor<i64>, tensor<i64>) -> tensor<2x2xi32>
// %result: [
//           [1, 1],
//           [1, 1]
//          ]

Ví dụ khác

dynamic_update_slice

Ngữ nghĩa

Tạo ra một tensor result bằng tensor operand, ngoại trừ việc lát cắt bắt đầu từ start_indices được cập nhật bằng các giá trị trong update. Chính thức hơn, result[result_index] được định nghĩa là:

  • update[update_index] nếu 0 <= update_index < shape(update), trong đó:
    • adjusted_start_indices = clamp(0, start_indices, shape(operand) - shape(update)).
    • update_index = result_index - adjusted_start_indices.
  • Nếu không, operand[result_index].

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C4), (C6)
(I2) update Tensor lượng tử hoá tensor hoặc mỗi tensor (C2), (C3), (C6)
(I3) start_indices số biến thiên của tensor 0 chiều của kiểu số nguyên (C4), (C5)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Giới hạn

  • (C1) type(operand) = type(result).
  • (C2) element_type(update) = element_type(operand).
  • (C3) rank(update) = rank(operand).
  • (C4) size(start_indices) = rank(operand).
  • (C5) same(type(start_indices...)).
  • (C6) 0 <= shape(update) <= shape(operand).

Ví dụ

// %operand: [
//            [1, 1, 0, 0],
//            [1, 1, 0, 0],
//            [1, 1, 1, 1],
//            [1, 1, 1, 1]
//           ]
// %update: [
//           [1, 1],
//           [1, 1]
//          ]
// %start_indices0: -1
// %start_indices1: 3
%result = "stablehlo.dynamic_update_slice"(%operand, %update, %start_indices0, %start_indices1)
  : (tensor<4x4xi32>, tensor<2x2xi32>, tensor<i64>, tensor<i64>) -> tensor<4x4xi32>
// %result: [
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1],
//           [1, 1, 1, 1]
//          ]

Ví dụ khác

hàm mũ

Ngữ nghĩa

Thực hiện phép toán luỹ thừa theo phần tử trên tensor operand và 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ố thực: exp từ IEEE-754.
  • Đối với số phức: số mũ phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(exponential, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.exponential"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[1.0, 2.7182818284590451], [7.3890560989306504, 20.085536923187668]]

Ví dụ khác

exponential_minus_one

Ngữ nghĩa

Thực hiện luỹ thừa phần tử theo cấp số nhân trừ đi một phép toá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 làm như sau:

  • Đối với số thực: expm1 từ IEEE-754.
  • Đối với số phức: số mũ phức trừ đi một.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(exponential_minus_one, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [0.0, 1.0]
%result = "stablehlo.exponential_minus_one"(%operand) : (tensor<2xf64>) -> tensor<2xf64>
// %result: [0.0, 1.71828187]

Ví dụ khác

fft

Ngữ nghĩa

Thực hiện các phép biến đổi Fourier tiến và nghịch đảo cho các dữ liệu đầu vào/đầu ra thực và phức tạp.

fft_type có một trong những trạng thái sau đây:

  • FFT: Chuyển tiếp FFT từ phức tạp đến phức tạp.
  • IFFT: FFT nghịch đảo từ phức tạp đến phức tạp.
  • RFFT: Chuyển tiếp FFT từ thực đến phức tạp.
  • IRFFT: nghịch đảo FFT từ thực đến phức (tức là lấy phức tạp, trả về thực).

Chính thức hơn, vì hàm fft lấy tensor 1 chiều của các loại phức tạp làm đầu vào, tạo ra tensor 1 chiều cùng loại với đầu ra và tính toán biến đổi Fourier riêng biệt:

Đối với fft_type = FFT, result được định nghĩa là kết quả cuối cùng của một chuỗi phép tính L, trong đó L = size(fft_length). Ví dụ: đối với L = 3:

  • result1[i0, ..., :] = fft(operand[i0, ..., :]).
  • result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]).

Hơn nữa, do hàm ifft có cùng chữ ký loại và tính toán nghịch đảo của fft:

Đối với fft_type = IFFT, result được định nghĩa là phép tính nghịch đảo của các phép tính cho fft_type = FFT. Ví dụ: đối với L = 3:

  • result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]).
  • result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :] = ifft(result2[i0, ..., :]).

Hơn nữa, do hàm rfft sử dụng tensor 1 chiều của các loại dấu phẩy động, tạo ra tensor 1 chiều của các loại phức tạp của cùng một ngữ nghĩa dấu phẩy động và hoạt động như sau:

  • rfft(real_operand) = truncated_result trong đó
  • complex_operand... = (real_operand..., 0.0).
  • complex_result = fft(complex_operand).
  • truncated_result = complex_result[:(rank(complex_result) / 2 + 1)].

(Khi tính toán phép biến đổi Fourier riêng biệt cho các toán hạng thực, phần tử N/2 + 1 đầu tiên của kết quả sẽ xác định rõ ràng phần còn lại của kết quả, vì vậy, kết quả của rfft sẽ được cắt bớt để tránh tính toán các phần tử thừa).

Đối với fft_type = RFFT, result được định nghĩa là kết quả cuối cùng của một chuỗi phép tính L, trong đó L = size(fft_length). Ví dụ: đối với L = 3:

  • result1[i0, ..., :] = rfft(operand[i0, ..., :]).
  • result2[i0, ..., :, iR-1] = fft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :, iR-2, iR-1] = fft(result2[i0, ..., :, iR-2, iR-1]).

Cuối cùng, hãy đưa ra hàm irfft có cùng chữ ký loại và tính toán nghịch đảo của rfft:

Đối với fft_type = IRFFT, result được định nghĩa là phép tính nghịch đảo của các phép tính cho fft_type = RFFT. Ví dụ: đối với L = 3:

  • result1[i0, ..., :, iR-2, iR-1] = ifft(operand[i0, ..., :, iR-2, iR-1]).
  • result2[i0, ..., :, iR-1] = ifft(result1[i0, ..., :, iR-1]).
  • result[i0, ..., :] = irfft(result2[i0, ..., :]).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của kiểu dấu phẩy động hoặc phức (C1), (C2), (C4), (C5)
(I2) fft_type enum của FFT, IFFT, RFFTIRFFT (C2), (C5)
(I3) fft_length Hằng số tensor 1 chiều thuộc loại si64 (C1), (C3), (C4)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu dấu phẩy động hoặc phức (C2), (C4), (C5)

Giới hạn

  • (C1) size(fft_length) <= rank(operand).
  • (C2) Mối quan hệ giữa các loại phần tử operandresult là khác nhau:
    • Nếu fft_type = FFT, element_type(operand)element_type(result) có cùng kiểu phức tạp.
    • Nếu fft_type = IFFT, element_type(operand)element_type(result) có cùng kiểu phức tạp.
    • Nếu fft_type = RFFT, element_type(operand) là một kiểu dấu phẩy động và element_type(result) là một loại phức tạp của cùng một ngữ nghĩa dấu phẩy động.
    • Nếu fft_type = IRFFT, element_type(operand) sẽ là một kiểu phức tạp và element_type(result) là một kiểu dấu phẩy động của cùng ngữ nghĩa dấu phẩy động.
  • (C3) 1 <= size(fft_length) <= 3.
  • (C4) Nếu trong số operandresult, có một tensor real thuộc kiểu dấu phẩy động, thì shape(real)[-size(fft_length):] = fft_length.
  • (C5) shape(result) = shape(operand) ngoại trừ:
    • Nếu fft_type = RFFT, dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1.
    • Nếu fft_type = IRFFT, dim(operand, -1) = dim(result, -1) = 0 ? 0 : dim(result, -1) / 2 + 1.

Ví dụ

// %operand: [(1.0, 0.0), (0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
%result = "stablehlo.fft"(%operand) {
  fft_type = #stablehlo<fft_type FFT>,
  fft_length = dense<4> : tensor<1xi64>
} : (tensor<4xcomplex<f32>>) -> tensor<4xcomplex<f32>>
// %result: [(1.0, 0.0), (1.0, 0.0), (1.0, 0.0), (1.0, 0.0)]

phần nguyên gần nhất bên trái

Ngữ nghĩa

Thực hiện tầng trên các phần tử của tensor operand và tạo ra một tensor result. Triển khai thao tác roundToIntegralTowardNegative từ quy cách IEEE-754. Đối với các loại lượng tử hoá, thực hiện dequantize_op_quantize(floor, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

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.floor"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-1.0, -1.0, 0.0, 0.0, 2.0]

Ví dụ khác

thu thập

Ngữ nghĩa

Thu thập các lát cắt từ tensor operand từ các độ lệch được chỉ định trong start_indices và tạo ra một tensor result.

Sơ đồ sau đây cho thấy cách các phần tử trong result ánh xạ với các phần tử trong operand bằng một ví dụ cụ thể. Sơ đồ này chọn một vài ví dụ về chỉ mục result và giải thích chi tiết các chỉ mục operand tương ứng.

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

  • batch_dims = [d for d in axes(result) and d not in offset_dims].
  • batch_index = result_index[batch_dims...].
  • start_index được khai báo là:
    • start_indices[bi0, ..., :, ..., biN], trong đó bi là các phần tử riêng lẻ trong batch_index: được chèn vào chỉ mục index_vector_dim, nếu index_vector_dim < rank(start_indices).
    • Nếu không, [start_indices[batch_index]].
  • Đối với d_operand trong axes(operand),
    • full_start_index[d_operand] = clamp(start_index[d_start], 0, dim(operand, d_operand) - slice_sizes[d_operand]) nếu d_operand = start_index_map[d_start].
    • Nếu không, full_start_index[d_operand] = 0.
  • offset_index = result_index[offset_dims...].
  • full_offset_index = [oi0, ..., 0, ..., oiN], trong đó oi là các phần tử riêng lẻ trong offset_index0 được chèn tại các chỉ mục từ collapsed_slice_dims.
  • operand_index = full_start_index + full_offset_index.

Nếu indices_are_sortedtrue thì quá trình triển khai có thể giả định rằng start_indices được sắp xếp theo start_index_map, nếu không thì hành vi sẽ không được xác định. Chính thức hơn, đối với tất cả i1 < i2 từ indices(result), full_start_index(i1) <= full_start_index(i2).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C7), (C10-C12), (C14)
(I2) start_indices tensor của kiểu số nguyên (C2), (C3), (C13)
(I3) offset_dims Hằng số tensor 1 chiều thuộc loại si64 (C1), (C4-C5), (C13)
(I4) collapsed_slice_dims Hằng số tensor 1 chiều thuộc loại si64 (C1), (C6-C8), (C13)
(I5) start_index_map Hằng số tensor 1 chiều thuộc loại si64 (C3), (C9), (C10)
(I6) index_vector_dim hằng số có kiểu si64 (C2), (C3), (C13)
(I7) slice_sizes Hằng số tensor 1 chiều thuộc loại si64 (C8), (C11-C13)
(I8) indices_are_sorted hằng số có kiểu i1

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C5), (C13-C14)

Giới hạn

  • (C1) rank(operand) = size(offset_dims) + size(collapsed_slice_dims).
  • (C2) 0 <= index_vector_dim <= rank(start_indices).
  • (C3) size(start_index_map) = index_vector_dim < rank(start_indices) ? dim(start_indices, index_vector_dim) : 1.
  • (C4) is_unique(offset_dims) and is_sorted(offset_dims).
  • (C5) 0 <= offset_dims < rank(result).
  • (C6) is_unique(collapsed_slice_dims) and is_sorted(collapsed_slice_dims).
  • (C7) 0 <= collapsed_slice_dims < rank(operand).
  • (C8) slice_sizes[collapsed_slice_dims...] <= 1.
  • (C9) is_unique(start_index_map).
  • (C10) 0 <= start_index_map < rank(operand).
  • (C11) size(slice_sizes) = rank(operand).
  • (C12) 0 <= slice_sizes <= shape(operand).
  • (C13) shape(result) = combine(batch_dim_sizes, offset_dim_sizes) trong đó:
    • batch_dim_sizes = shape(start_indices) ngoại trừ kích thước kích thước của start_indices tương ứng với index_vector_dim.
    • offset_dim_sizes = shape(slice_sizes) ngoại trừ việc kích thước kích thước trong slice_sizes tương ứng với collapsed_slice_dims sẽ không được đưa vào.
    • combine đặt batch_dim_sizes ở các trục tương ứng với batch_dimsoffset_dim_sizes tại các trục tương ứng với offset_dims.
  • (C14) element_type(operand) = element_type(result).

Ví dụ

// %operand: [
//            [[1, 2], [3, 4], [5, 6], [7, 8]],
//            [[9, 10],[11, 12], [13, 14], [15, 16]],
//            [[17, 18], [19, 20], [21, 22], [23, 24]]
//           ]
// %start_indices: [
//                  [[0, 0], [1, 0], [2, 1]],
//                  [[0, 1], [1, 1], [0, 2]]
//                 ]
%result = "stablehlo.gather"(%operand, %start_indices) {
  dimension_numbers = #stablehlo.gather<
    offset_dims = [2, 3],
    collapsed_slice_dims = [0],
    start_index_map = [1, 0],
    index_vector_dim = 2>,
  slice_sizes = dense<[1, 2, 2]> : tensor<3xi64>,
  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
// %result: [
//            [
//              [[1, 2], [3, 4]],
//              [[3, 4], [5, 6]],
//              [[13, 14], [15, 16]]
//            ],
//            [
//              [[9, 10], [11, 12]],
//              [[11, 12], [13, 14]],
//              [[17, 18], [19, 20]]
//            ]
//          ]

Ví dụ khác

get_dimension_size

Ngữ nghĩa

Tạo kích thước của dimension đã cho của operand. Chính thức hơn là result = dim(operand, dimension).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor (C1)
(I2) dimension hằng số có kiểu si64 (C1)

Kết quả đầu ra

Tên Loại
result Tensor 0 chiều thuộc loại si32

Giới hạn

  • (C1) 0 <= dimension < rank(operand).

Ví dụ

// %operand: [[1, 2, 3], [4, 5, 6]]
%result = "stablehlo.get_dimension_size"(%operand) {
  dimension = 1 : i64
} : (tensor<2x3xi64>) -> tensor<i32>
// %result: 3

Ví dụ khác

get_tuple_element

Ngữ nghĩa

Trích xuất phần tử tại vị trí index của bộ dữ liệu operand và tạo ra một result. Chính thức hơn là result = operand[index].

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tuple (C1), (C2)
(I2) index hằng số có kiểu si32 (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result mọi loại được hỗ trợ (C2)

Giới hạn

  • (C1) 0 <= index < size(operand).
  • (C2) type(result) = tuple_element_types(operand)[index].

Ví dụ

// %operand: ([1.0, 2.0], (3))
%result = "stablehlo.get_tuple_element"(%operand) {
  index = 0 : i32
} : (tuple<tensor<2xf32>, tuple<tensor<i32>>>) -> tensor<2xf32>
// %result: [1.0, 2.0]

Ví dụ khác

if

Ngữ nghĩa

Tạo đầu ra từ việc thực thi chính xác một hàm từ true_branch hoặc false_branch tuỳ thuộc vào giá trị của pred. Chính thức hơn là result = pred ? true_branch() : false_branch().

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) pred Tensor 0 chiều thuộc loại i1
(I2) true_branch hàm (C1–C3)
(I3) false_branch hàm (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên, tensor lượng tử hoá hoặc mã thông báo (C3)

Giới hạn

  • (C1) input_types(true_branch) = input_types(false_branch) = [].
  • (C2) output_types(true_branch) = output_types(false_branch).
  • (C3) type(results...) = output_types(true_branch).

Ví dụ

// %result_true_branch: 10
// %result_false_branch: 11
// %pred: true
%result = "stablehlo.if"(%pred) ({
  "stablehlo.return"(%result_true_branch) : (tensor<i32>) -> ()
}, {
  "stablehlo.return"(%result_false_branch) : (tensor<i32>) -> ()
}) : (tensor<i1>) -> tensor<i32>
// %result: 10

Ví dụ khác

hình ảnh

Ngữ nghĩa

Trích xuất phần ảo từ operand và tạo ra tensor result. Chính thức hơn, đối với mỗi phần tử x: imag(x) = is_complex(x) ? imaginary_part(x) : constant(0, element_type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của kiểu dấu phẩy động hoặc phức (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động (C1), (C2)

Giới hạn

  • (C1) shape(result) = shape(operand).
  • (C2) element_type(result) được định nghĩa là:
    • complex_element_type(element_type(operand)) nếu is_complex(operand).
    • Nếu không, element_type(operand).

Ví dụ

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.imag"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [2.0, 4.0]

Ví dụ khác

nguồn cấp dữ liệu

Ngữ nghĩa

Đọc dữ liệu từ nguồn cấp dữ liệu và tạo ra results.

Ngữ nghĩa của infeed_config được xác định dựa trên phương thức triển khai.

results bao gồm các giá trị tải trọng xuất hiện trước và một mã thông báo xuất hiện sau cùng. Trong tương lai, chúng tôi dự định chia tải trọng và mã thông báo thành hai dữ liệu đầu ra riêng biệt để cải thiện tính rõ ràng (#670).

Thông tin đầu vào

Nhãn Tên Loại
(I1) token token
(I2) infeed_config hằng số có kiểu string

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên, tensor lượng tử hoá hoặc mã thông báo (C1–C3)

Giới hạn

  • (C1) 0 < size(results).
  • (C2) is_empty(result[:-1]) hoặc is_tensor(type(results[:-1])).
  • (C3) is_token(type(results[-1])).

Ví dụ

// %token: !stablehlo.token
// infeed_queue[0]: [[1, 2], [3, 4]]
// infeed_queue[1]: [[5, 6], [7, 8]]
%results0:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results0#0: [[1, 2], [3, 4]]
%results1:2 = "stablehlo.infeed"(%token) {
  infeed_config = ""
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)
// results1#0: [[5, 6], [7, 8]]

Ví dụ khác

Iota

Ngữ nghĩa

Điền các giá trị vào tensor output theo thứ tự tăng dần bắt đầu từ 0 theo chiều iota_dimension. Chính thức hơn,

output[result_index] = constant(is_quantized(output) ? quantize(result_index[iota_dimension], element_type(output)) : result_index[iota_dimension], element_type(output)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) iota_dimension si64 (C1)

Kết quả đầu ra

Tên Loại Giới hạn
output tensor của số nguyên, dấu phẩy động hoặc loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

  • (C1) 0 <= iota_dimension < rank(output).

Ví dụ

%output = "stablehlo.iota"() {
  iota_dimension = 0 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 0, 0, 0, 0],
//           [1, 1, 1, 1, 1],
//           [2, 2, 2, 2, 2],
//           [3, 3, 3, 3, 3]
//          ]

%output = "stablehlo.iota"() {
  iota_dimension = 1 : i64
} : () -> tensor<4x5xi32>
// %output: [
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4],
//           [0, 1, 2, 3, 4]
//          ]

Ví dụ khác

is_finite

Ngữ nghĩa

Thực hiện kiểm tra theo hướng phần tử xem giá trị trong x có hữu hạn hay không (tức là không phải +Inf, -Inf hoặc NaN) và tạo ra tensor y. Triển khai toán tử isFinite từ thông số kỹ thuật IEEE-754. Đối với các loại lượng tử hoá, kết quả luôn là true.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) x tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
y tensor của loại boolean (C1)

Giới hạn

  • (C1) shape(x) = shape(y).

Ví dụ

// Logical values: -Inf, +Inf, NaN, ...
// %x: [0xFFF0000000000000, 0x7FF0000000000000, 0x7FF8000000000000, -10.0, -0.0, 0.0, 10.0]
%y = "stablehlo.is_finite"(%x) : (tensor<7xf64) -> tensor<7xi1>
// %y: [false, false, false, true, true, true, true]

Ví dụ khác

log

Ngữ nghĩa

Thực hiện phép toán lôgarit dựa trên 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: log từ IEEE-754.
  • Đối với số phức: lôgarit phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(log, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [[1.0, 2.0], [3.0, 4.0]]
%result = "stablehlo.log"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.0, 0.69314718055994529], [1.0986122886681098, 1.3862943611198906]]

Ví dụ khác

log_plus_one

Ngữ nghĩa

Thực hiện lôgarit theo phần tử cộng với một phép toá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 làm như sau:

  • Đối với số thực: logp1 từ IEEE-754.
  • Đối với số phức: lôgarit phức cộng với 1.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(log_plus_one, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [0.0, -0.999, 7.0, 6.38905621, 15.0]
%result = "stablehlo.log_plus_one"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [0.0, -6.90776825, 2.07944155, 2.0, 2.77258873]

Ví dụ khác

hậu cần

Ngữ nghĩa

Thực hiện thao tác hậu cần tương ứng với các 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: division(1, addition(1, exp(-x))) từ IEEE-754.
  • Đối với số phức: logistic phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(logistic, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [[0.0, 1.0], [2.0, 3.0]]
%result = "stablehlo.logistic"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
// %result: [[0.5, 0.73105858], [0.88079708, 0.95257413]]

Ví dụ khác

map

Ngữ nghĩa

Áp dụng hàm ánh xạ computation cho inputs cùng với dimensions và tạo tensor result.

Chính thức hơn là result[result_index] = computation(inputs...[result_index]). Xin lưu ý rằng dimensions hiện không được dùng và có thể sẽ bị xoá trong tương lai (#487).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1–C4)
(I2) dimensions Hằng số tensor 1 chiều thuộc loại si64 (C3)
(I3) computation hàm (C4)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C4)

Giới hạn

  • (C1) shape(inputs...) = shape(result).
  • (C2) 0 < size(inputs) = N.
  • (C3) dimensions = range(rank(inputs[0])).
  • (C4) computation có loại (tensor<E0>, ..., tensor<EN-1>) -> tensor<E'>, trong đó Ei = element_type(inputs[i])E' = element_type(result).

Ví dụ

// %input0: [[0, 1], [2, 3]]
// %input1: [[4, 5], [6, 7]]
%result = "stablehlo.map"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<i64>
    stablehlo.return %0 : tensor<i64>
}) {
  dimensions = dense<[0, 1]> : tensor<2xi64>
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> tensor<2x2xi64>
// %result: [[0, 5], [12, 21]]

Ví dụ khác

tối đa

Ngữ nghĩa

Thực hiện thao tác tối đa trên tensor lhsrhs, đồ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 boolean: toán tử logic OR.
  • Đối với số nguyên: số nguyên tối đa.
  • Đối với số thực: maximum từ IEEE-754.
  • Đối với số phức: giá trị tối đa trong từ vựng cho cặp (real, imaginary). Việc áp đặt thứ tự trên các số phức liên quan đến ngữ nghĩa bất ngờ, vì vậy, trong tương lai, chúng tôi dự định ngừng hỗ trợ số phức cho toán tử này (#560).
  • Đối với các loại lượng tử hoá:
    • dequantize_op_quantize(maximum, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)
(I2) rhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Giới hạn

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

Ví dụ

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

Ví dụ khác

tối thiểu

Ngữ nghĩa

Thực hiện thao tác tối thiểu theo phần tử trên tensor lhsrhs, đồ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 boolean: logic AND.
  • Đối với số nguyên: số nguyên tối thiểu.
  • Đối với số thực: minimum từ IEEE-754.
  • Đối với số phức: từ vựng tối thiểu cho cặp (real, imaginary). Việc áp đặt thứ tự trên các số phức liên quan đến ngữ nghĩa bất ngờ, vì vậy, trong tương lai, chúng tôi dự định ngừng hỗ trợ số phức cho toán tử này (#560).
  • Đối với các loại lượng tử hoá:
    • dequantize_op_quantize(minimum, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)
(I2) rhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Giới hạn

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

Ví dụ

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

Ví dụ khác

nhân

Ngữ nghĩa

Thực hiện tích số phần tử của hai tensor lhsrhs, đồ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 AND.
  • Đối với số nguyên: phép nhân số nguyên.
  • Đối với số thực: multiplication từ IEEE-754.
  • Đối với số phức: phép nhân phức.
  • Đối với các loại lượng tử hoá:
    • dequantize_op_quantize(multiply, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)
(I2) rhs Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1)

Giới hạn

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

Ví dụ

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

Ví dụ khác

phủ định

Ngữ nghĩa

Thực hiện phủ định phần tử của 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ố nguyên có dấu: số nguyên phủ định.
  • Đối với số nguyên chưa ký: truyền bit về số nguyên đã ký, giá trị phủ định số nguyên, truyền bit trở lại số nguyên chưa ký.
  • Đối với số thực: negate từ IEEE-754.
  • Đối với số phức: dạng phủ định phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(negate, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// Negation operation with integer Tensors
// %operand: [0, -2]
%result = "stablehlo.negate"(%operand) : (tensor<2xi32>) -> tensor<2xi32>
// %result: [0, 2]

// Negation operation with with complex tensors
// %operand: (2.5, 0.0)
%result = "stablehlo.negate"(%operand) : (tensor<1xcomplex<f32>>) -> tensor<1xcomplex<f32>>
// %result: [-2.5, -0.0]

Ví dụ khác

không

Ngữ nghĩa

Thực hiện KHÔNG PHẢI của tensor operand theo phần tử 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 boolean: logic NOT.
  • Đối với số nguyên: bitwise NOT.

Đối số

Tên Loại Giới hạn
operand tensor của kiểu boolean hoặc số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu boolean hoặc số nguyên (C1)

Giới hạn

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

Ví dụ

// Bitwise operation with with integer tensors
// %operand: [[1, 2], [3, 4]]
%result = "stablehlo.not"(%operand) : (tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[-2, -3], [-4, -5]]

// Bitwise operation with with boolean tensors
// %operand: [true, false]
%result = "stablehlo.not"(%operand) : (tensor<2xi1>) -> tensor<2xi1>
// %result: [false, true]

optimization_barrier

Ngữ nghĩa

Đảm bảo rằng các thao tác tạo ra operand được thực thi trước bất kỳ thao tác nào phụ thuộc vào result và ngăn việc chuyển đổi trình biên dịch không di chuyển các hoạt động qua rào cản. Ngoài ra, thao tác này là một mã nhận dạng, tức là result = operand.

Đối số

Tên Loại Giới hạn
operand số lượng tensor biến thiên, tensor hoặc mã thông báo lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result số lượng tensor biến thiên, tensor hoặc mã thông báo lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand0: 0.0
// %operand1: 1.0
%result0, %result1 = "stablehlo.optimization_barrier"(%operand0, %operand1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
// %result0: 0.0
// %result1: 1.0

Ví dụ khác

hoặc

Ngữ nghĩa

Thực hiện OR phần tử OR của hai tensor lhsrhs, đồ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: toán tử logic OR.
  • Đối với số nguyên: bitwise OR.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của kiểu số nguyên hoặc boolean (C1)
(I2) rhs tensor của kiểu số nguyên hoặc boolean (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên hoặc boolean (C1)

Giới hạn

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

Ví dụ

// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 6], [7, 12]]

// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.or"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, true]]

bên ngoài

Ngữ nghĩa

Ghi inputs vào nguồn cấp dữ liệu bên ngoài và tạo mã thông báo result.

Ngữ nghĩa của outfeed_config được xác định dựa trên phương thức triển khai.

Thông tin đầu vào

Nhãn Tên Loại
(I1) inputs số variadic của tensor hoặc tensor lượng tử hoá
(I2) token token
(I3) outfeed_config hằng số có kiểu string

Kết quả đầu ra

Tên Loại
result token

Ví dụ

%result = "stablehlo.outfeed"(%inputs0, %token) {
  outfeed_config = ""
} : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token

Ví dụ khác

bàn phím

Ngữ nghĩa

Mở rộng operand bằng khoảng đệm xung quanh tensor cũng như giữa các phần tử của tensor với padding_value đã cho.

edge_padding_lowedge_padding_high chỉ định số lượng khoảng đệm đã thêm ở cấp thấp (bên cạnh chỉ mục 0) và cao cấp (bên cạnh chỉ mục cao nhất) của từng chiều tương ứng. Số lượng khoảng đệm có thể là số âm, trong đó giá trị tuyệt đối của khoảng đệm âm cho biết số phần tử cần xoá khỏi kích thước đã chỉ định.

interior_padding chỉ định khoảng đệm được thêm vào giữa hai phần tử bất kỳ trong mỗi chiều (không được là số âm). Khoảng đệm bên trong xuất hiện trước khoảng đệm cạnh, như vậy khoảng đệm cạnh âm sẽ xoá các phần tử khỏi toán hạng được đệm bên trong.

Chính thức hơn, result[result_index] được định nghĩa là:

  • operand[operand_index] nếu result_index = edge_padding_low + operand_index * (interior_padding + 1).
  • Nếu không, padding_value.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C2), (C4)
(I2) padding_value Tensor lượng tử hoá tensor 0 chiều hoặc mỗi tensor (C1)
(I3) edge_padding_low Hằng số tensor 1 chiều thuộc loại si64 (C1), (C4)
(I4) edge_padding_high Hằng số tensor 1 chiều thuộc loại si64 (C1), (C4)
(I5) interior_padding Hằng số tensor 1 chiều thuộc loại si64 (C2–C4)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C3–C6)

Giới hạn

  • (C1) element_type(operand) = element_type(padding_value) = element_type(result).
  • (C2) size(edge_padding_low) = size(edge_padding_high) = size(interior_padding) = rank(operand).
  • (C3) 0 <= interior_padding.
  • (C4) shape(result) = shape(operand) + edge_padding_low + max(shape(operand) - 1, 0) * interior_padding + edge_padding_high.

Ví dụ

// %operand: [
//            [1, 2, 3],
//            [4, 5, 6]
//           ]
// %padding_value: 0
%result = "stablehlo.pad"(%operand, %padding_value) {
  edge_padding_low = dense<[0, 1]> : tensor<2xi64>,
  edge_padding_high = dense<[2, 1]> : tensor<2xi64>,
  interior_padding = dense<[1, 2]> : tensor<2xi64>
} : (tensor<2x3xi32>, tensor<i32>) -> tensor<5x9xi32>
// %result: [
//           [0, 1, 0, 0, 2, 0, 0, 3, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 4, 0, 0, 5, 0, 0, 6, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0],
//           [0, 0, 0, 0, 0, 0, 0, 0, 0]
//          ]

Ví dụ khác

partition_id

Ngữ nghĩa

Tạo ra partition_id của quá trình hiện tại.

Kết quả đầu ra

Tên Loại
result Tensor 0 chiều thuộc loại ui32

Ví dụ

%result = "stablehlo.partition_id"() : () -> tensor<ui32>

Ví dụ khác

hoa popcnt

Ngữ nghĩa

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

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của kiểu số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên (C1)

Giới hạn

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

Ví dụ

// %operand: [0, 1, 2, 127]
%result = "stablehlo.popcnt"(%operand) : (tensor<4xi64>) -> tensor<4xi64>
// %result: [0, 1, 1, 7]

Ví dụ khác

công suất

Ngữ nghĩa

Thực hiện luỹ thừa phần tử của tensor lhs bằng tensor rhs 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ố nguyên: phép luỹ thừa số nguyên.
  • Đối với số thực: pow từ IEEE-754.
  • Đối với số phức: phép luỹ thừa phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(power, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) rhs tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %lhs: [-2.0, -0.0, -36.0, 5.0, 3.0, 10000.0]
// %rhs: [2.0, 2.0, 1.1, 2.0, -1.0, 10.0]
%result = "stablehlo.power"(%lhs, %rhs) : (tensor<6xf64>, tensor<6xf64>) -> tensor<6xf64>
// %result: [4.0, 0.0, -nan, 25.0, 0.333333343, inf]

Ví dụ khác

thực

Ngữ nghĩa

Trích xuất phần thực, theo phần tử, từ operand và tạo ra một tensor result. Chính thức hơn, đối với mỗi phần tử x: real(x) = is_complex(x) ? real_part(x) : x.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của kiểu dấu phẩy động hoặc phức (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động (C1), (C2)

Giới hạn

  • (C1) shape(result) = shape(operand).
  • (C2) element_type(result) được định nghĩa là:
    • complex_element_type(element_type(operand)) nếu is_complex(operand).
    • Nếu không, element_type(operand).

Ví dụ

// %operand: [(1.0, 2.0), (3.0, 4.0)]
%result = "stablehlo.real"(%operand) : (tensor<2xcomplex<f32>>) -> tensor<2xf32>
// %result: [1.0, 3.0]

Ví dụ khác

phản hồi

Ngữ nghĩa

Nhận dữ liệu từ một kênh bằng channel_id và tạo results.

Nếu is_host_transfertrue, thì thao tác này sẽ chuyển dữ liệu từ máy chủ. Nếu không, thiết bị sẽ chuyển dữ liệu từ một thiết bị khác. Điều này có nghĩa là việc triển khai được xác định. Cờ này sao chép thông tin được cung cấp trong channel_type, vì vậy, trong tương lai, chúng tôi dự định chỉ giữ lại một trong số đó (#666).

results bao gồm các giá trị tải trọng xuất hiện trước và một mã thông báo xuất hiện sau cùng. Trong tương lai, chúng tôi dự định chia tải trọng và mã thông báo thành hai dữ liệu đầu ra riêng biệt để cải thiện tính rõ ràng (#670).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) token token (C4)
(I2) channel_id hằng số có kiểu si64
(I3) channel_type enum của DEVICE_TO_DEVICEHOST_TO_DEVICE (C1)
(I4) is_host_transfer hằng số có kiểu i1 (C1)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên, tensor lượng tử hoá hoặc mã thông báo (C2–C4)

Giới hạn

  • (C1) channel_type được định nghĩa là:
    • HOST_TO_DEVICE nếu is_host_transfer = true,
    • Nếu không, DEVICE_TO_DEVICE.
  • (C2) 0 < size(results).
  • (C3) is_empty(result[:-1]) hoặc is_tensor(type(results[:-1])).
  • (C4) is_token(type(results[-1])).

Ví dụ

%results0, %results1 = "stablehlo.recv"(%token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
  is_host_transfer = true
} : (!stablehlo.token) -> (tensor<2x2xi64>, !stablehlo.token)

Ví dụ khác

giảm

Ngữ nghĩa

Áp dụng hàm rút gọn body cho inputsinit_values cùng với dimensions và tạo ra tensor results.

Thứ tự rút gọn được xác định theo phương thức triển khai, có nghĩa là bodyinit_values phải tạo thành một đơn thể (monoid) để đảm bảo thao tác này tạo ra cùng một kết quả cho tất cả dữ liệu đầu vào trong tất cả các phương thức triển khai. Tuy nhiên, điều kiện này không có nhiều mức giảm phổ biến. Ví dụ: phép thêm dấu phẩy động cho body và 0 cho init_values không thực sự tạo thành một đơn sắc vì việc bổ sung dấu phẩy động không liên kết.

Chính thức hơn, results...[j0, ..., jR-1] = reduce(input_slices_converted), trong đó:

  • input_slices = inputs...[j0, ..., :, ..., jR-1], trong đó : được chèn tại dimensions.
  • input_slices_converted = to_destination_type(input_slices..., type(func_inputs(body)[:len(func_inputs(body))//2])...).
  • init_values_converted = to_destination_type(init_values..., type(func_inputs(body)[len(func_inputs(body))//2:])...).
  • reduce(input_slices_converted) = exec(schedule) cho một số cây nhị phân schedule, trong đó:
    • exec(node) = body(exec(node.left), exec(node.right)).
    • exec(leaf) = leaf.value.
  • schedule là một cây nhị phân đầy đủ do triển khai xác định, trong đó truyền tải theo thứ tự bao gồm:
    • Giá trị input_slices_converted...[index], cho tất cả index trong index_space(input_slices_converted) theo thứ tự từ vựng tăng dần của index.
    • Xen kẽ với một lượng init_values_converted do quá trình triển khai xác định tại các vị trí do phương thức triển khai xác định.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1-C4), (C6), (C7)
(I2) init_values số biến thiên của tensor 0 chiều hoặc tensor lượng tử hoá mỗi tensor (C2), (C3)
(I3) dimensions Hằng số tensor 1 chiều thuộc loại si64 (C4), (C5), (C7)
(I4) body hàm (C6)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C3), (C7), (C8)

Giới hạn

  • (C1) same(shape(inputs...)).
  • (C2) element_type(inputs...) = element_type(init_values...).
  • (C3) 0 < size(inputs) = size(init_values) = size(results) = N.
  • (C4) 0 <= dimensions < rank(inputs[0]).
  • (C5) is_unique(dimensions).
  • (C6) body có loại (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>), trong đó is_promotable(element_type(inputs[i]), Ei).
  • (C7) shape(results...) = shape(inputs...) ngoại trừ việc kích thước kích thước của inputs... tương ứng với dimensions sẽ không được đưa vào.
  • (C8) element_type(results[i]) = Ei cho tất cả i trong [0,N).

Ví dụ

// %input = [[0, 1, 2, 3, 4, 5]]
// %init_value = 0
%result = "stablehlo.reduce"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  dimensions = dense<1> : tensor<1xi64>
} : (tensor<1x6xi64>, tensor<i64>) -> tensor<1xi64>
// %result = [15]

Ví dụ khác

reduce_precision

Ngữ nghĩa

Thực hiện việc chuyển đổi operand thành phần tử theo kiểu dấu phẩy động khác sử dụng exponent_bitsmantissa_bits, quay lại kiểu dấu phẩy động ban đầu, đồng thời tạo tensor output.

Chính thức hơn:

  • Các bit mantissa của giá trị gốc được cập nhật để làm tròn giá trị gốc đến giá trị gần nhất biểu thị bằng mantissa_bits bằng cách sử dụng ngữ nghĩa roundToIntegralTiesToEven.
  • Sau đó, nếu mantissa_bits nhỏ hơn số bit mantissa của giá trị ban đầu, thì các bit mantissa sẽ bị cắt bớt thành mantissa_bits.
  • Sau đó, nếu các bit số mũ của kết quả trung gian không phù hợp với phạm vi do exponent_bits cung cấp, thì kết quả trung gian sẽ tràn đến vô hạn bằng cách sử dụng dấu ban đầu hoặc luồng dưới về 0 khi sử dụng dấu gốc.
  • Đối với các loại lượng tử hoá, thực hiện dequantize_op_quantize( lambda operand: reduce_precision(operand, exponent_bits, mantissa_bits), operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) exponent_bits hằng số có kiểu si32 (C2)
(I3) mantissa_bits hằng số có kiểu si32 (C3)

Kết quả đầu ra

Tên Loại Giới hạn
output tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

  • (C1) baseline_type(operand) = baseline_type(output).
  • (C2) 1 <= exponent_bits.
  • (C3) 0 <= mantissa_bits.

Ví dụ

// Logical values: +Inf, NaN, +Denormal, 0.0, 65519.0, 65520.0
// %operand: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0x0000000000000001, 0.0, 65519.0, 65520.0]
%output = "stablehlo.reduce_precision"(%operand) {
  exponent_bits = 5 : i32,
  mantissa_bits = 10 : i32
} : (tensor<6xf64>) -> tensor<6xf64>
// Logical values: +Inf, NaN, 0.0, 0.0, 65504.0, +Inf
// %output: [0x7FF0000000000000, 0x7FFFFFFFFFFFFFFF, 0.0, 0.0, 65504.0, 0x7FF0000000000000]

Ví dụ khác

reduce_scatter

Ngữ nghĩa

Trong mỗi nhóm quy trình trong lưới quy trình StableHLO, thực hiện rút gọn bằng cách sử dụng computations trên các giá trị của tensor operand của mỗi quy trình, chia kết quả rút gọn dọc theo scatter_dimension thành nhiều phần và phân tán các phần đã tách giữa các quy trình để tạo result.

Thao tác này sẽ chia lưới quy trình StableHLO thành process_groups đượ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:

  • reduced_value = all_reduce(operand, replica_groups, channel_id, use_global_device_ids, computation).
  • parts@sender = split(reduced_value@sender, dim(process_groups, 1), scatter_dimension).
  • result@receiver = parts@sender[receiver_index] cho tất cả sender trong process_group, trong đó receiver_index = process_group.index(receiver).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C2), (C7), (C8)
(I2) scatter_dimension hằng số có kiểu si64 (C1), (C2), (C8)
(I3) replica_groups Hằng số tensor 2 chiều thuộc loại si64 (C3–C5)
(I4) channel_id hằng số có kiểu si64 (C6)
(I5) use_global_device_ids hằng số có kiểu i1 (C6)
(I6) computation hàm (C7)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C8–C9)

Giới hạn

  • (C1) dim(operand, scatter_dimension) % dim(process_groups, 1) = 0.
  • (C2) 0 <= scatter_dimension < rank(operand).
  • (C3) is_unique(replica_groups).
  • (C4) size(replica_groups) được định nghĩa là:
    • num_replicas nếu sử dụng cross_replica.
    • num_replicas nếu sử dụng cross_replica_and_partition.
    • num_processes nếu sử dụng flattened_ids.
  • (C5) 0 <= replica_groups < size(replica_groups).
  • (C6) Nếu là use_global_device_ids = true, thì channel_id > 0.
  • (C7) computation có loại (tensor<E>, tensor<E>) -> (tensor<E>), trong đó is_promotable(element_type(operand), E).
  • (C8) shape(result) = shape(operand) ngoại trừ:
    • dim(result, scatter_dimension) = dim(operand, scatter_dimension) / dim(process_groups, 1).
  • (C9) element_type(result) = E.

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand@(0, 0): [[1, 2, 3, 4],
//                   [5, 6, 7, 8]]
// %operand@(1, 0): [[9, 10, 11, 12],
//                   [13, 14, 15, 16]]
%result = "stablehlo.reduce_scatter"(%operand) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
  "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor<2x4xi64>) -> tensor<2x2xi64>
//
// %result@(0, 0): [[10, 12],
//                  [18, 20]]
// %result@(1, 0): [[14, 16],
//                  [22, 24]]

Ví dụ khác

reduce_window

Ngữ nghĩa

Áp dụng hàm rút gọn body cho cửa sổ inputsinit_values rồi tạo results.

Sơ đồ sau đây cho thấy cách tính toán các phần tử trong results... từ inputs..., bằng cách sử dụng một ví dụ cụ thể.

Chính thức hơn, results...[result_index] = reduce(windows, init_values, axes(inputs...), body) (xem giảm thiểu) trong đó:

  • padded_inputs = pad(inputs..., init_values..., padding[:, 0], padding[:, 1], base_dilations - 1).
  • window_start = result_index * window_strides.
  • window_end = window_start + (window_dimensions - 1) * window_dilations + 1.
  • windows = slice(padded_inputs..., window_start, window_end, window_dilations).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1–C4), (C6), (C8), (C10), (C12), (C13), (C15)
(I2) init_values số biến thiên của tensor 0 chiều hoặc tensor lượng tử hoá mỗi tensor (C1), (C13)
(I3) window_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C4), (C5), (C15)
(I4) window_strides Hằng số tensor 1 chiều thuộc loại si64 (C6), (C7), (C15)
(I5) base_dilations Hằng số tensor 1 chiều thuộc loại si64 (C8), (C9), (C15)
(I6) window_dilations Hằng số tensor 1 chiều thuộc loại si64 (C10), (C11), (C15)
(I7) padding Hằng số tensor 2 chiều thuộc loại si64 (C12), (C15)
(I8) body hàm (C13)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1), (C14-C16)

Giới hạn

  • (C1) 0 < size(inputs) = size(init_values) = size(results) = N.
  • (C2) same(shape(inputs...)).
  • (C3) element_type(inputs...) = element_type(init_values...).
  • (C4) size(window_dimensions) = rank(inputs[0]).
  • (C5) 0 < window_dimensions.
  • (C6) size(window_strides) = rank(inputs[0]).
  • (C7) 0 < window_strides.
  • (C8) size(base_dilations) = rank(inputs[0]).
  • (C9) 0 < base_dilations.
  • (C10) size(window_dilations) = rank(inputs[0]).
  • (C11) 0 < window_dilations.
  • (C12) shape(padding) = [rank(inputs[0]), 2].
  • (C13) body có loại (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>), trong đó is_promotable(element_type(inputs[i]), Ei).
  • (C14) same(shape(results...)).
  • (C15) shape(results[0]) = num_windows trong đó:
    • dilated_input_shape = shape(inputs[0]) = 0 ? 0 : (shape(inputs[0]) - 1) * base_dilations + 1.
    • padded_input_shape = padding[:, 0] + dilated_input_shape + padding[:, 1].
    • dilated_window_shape = (window_dimensions - 1) * window_dilations + 1.
    • is_empty_window = padded_input_shape = 0 || dilated_window_shape > padded_input_shape.
    • num_windows = is_empty_window ? 0 : floor((padded_input_shape - dilated_window_shape) / window_strides) + 1.
  • (C16) element_type(results[i]) = Ei cho tất cả i trong [0,N).

Ví dụ

// %input = [[1, 2], [3, 4], [5, 6]]
// %init_value = 0
%result = "stablehlo.reduce_window"(%input, %init_value) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  window_dimensions = dense<[2, 1]> : tensor<2xi64>,
  window_strides = dense<[4, 1]> : tensor<2xi64>,
  base_dilations = dense<[2, 1]> : tensor<2xi64>,
  window_dilations = dense<[3, 1]> : tensor<2xi64>,
  padding = dense<[[2, 1], [0, 0]]> : tensor<2x2xi64>
} : (tensor<3x2xi64>, tensor<i64>) -> tensor<2x2xi64>
// %result = [[0, 0], [3, 4]]

Ví dụ khác

phần còn lại

Ngữ nghĩa

Thực hiện phần còn lại theo phần tử của tensor bị chia lhs và số chia rhs, đồng thời tạo ra tensor result.

Nói một cách chính thức hơn, dấu của kết quả được lấy từ số bị chia và giá trị tuyệt đối của kết quả luôn nhỏ hơn giá trị tuyệt đối của số chia. Phần còn lại được tính bằng lhs - d * rhs, trong đó d được tính theo công thức sau:

  • Đối với số nguyên: stablehlo.divide(lhs, rhs).
  • Đối với độ chính xác đơn: division(lhs, rhs) từ IEEE-754 có thuộc tính làm tròn roundTowardZero.
  • Đối với số phức: TBD (#997).
  • Đối với các loại lượng tử hoá:
    • dequantize_op_quantize(remainder, lhs, rhs, type(result)).

Đối với các loại phần tử dấu phẩy động, toán tử này trái ngược với toán tử remainder từ thông số kỹ thuật IEEE-754, trong đó d là một giá trị tích phân gần nhất với giá trị chính xác của lhs/rhs với mối quan hệ chẵn.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của số nguyên, dấu phẩy động hoặc loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) rhs tensor của số nguyên, dấu phẩy động hoặc loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của số nguyên, dấu phẩy động hoặc loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %lhs: [17, -17, 17, -17]
// %rhs: [3, 3, -3, -3]
%result = "stablehlo.remainder"(%lhs, %rhs) : (tensor<4xi64>, tensor<4xi64>) -> tensor<4xi64>
// %result: [2, -2, 2, -2]

Ví dụ khác

replica_id

Ngữ nghĩa

Tạo ra replica_id của quá trình hiện tại.

Kết quả đầu ra

Tên Loại
result Tensor 0 chiều thuộc loại ui32

Ví dụ

%result = "stablehlo.replica_id"() : () -> tensor<ui32>

Ví dụ khác

đổi hình dạng

Ngữ nghĩa

Thực hiện đổi hình dạng tensor operand thành tensor result. Về mặt lý thuyết, bạn cần giữ nguyên cách biểu diễn chính tắc nhưng có thể thay đổi hình dạng, chẳng hạn như từ tensor<2x3xf32> thành tensor<3x2xf32> hoặc tensor<6xf32>.

Chính thức hơn, result[result_index] = operand[operand_index], trong đó result_indexoperand_index có cùng vị trí theo thứ tự từ vựng là index_space(result)index_space(operand).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor hoặc tensor lượng tử hoá (C1–C3)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor hoặc tensor lượng tử hoá (C1–C3)

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)quantization_dimension(result) có thể khác nhau, nếu không.
  • (C2) size(operand) = size(result).
  • (C3) Nếu is_per_axis_quantized(operand):
    • reduce(dims(operand, [0, 1, ..., quantization_dimension(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [0, 1, ..., quantization_dimension(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y).
    • dim(operand, quantization_dimension(operand)) = dim(result, quantization_dimension(result)).
    • reduce(dims(operand, [quantization_dimension(operand) + 1, ..., rank(operand) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y) = reduce(dims(result, [quantization_dimension(result) + 1, ..., rank(result) - 1]), init_values=1, dimensions=[0], body=lambda x, y: x * y).

Ví dụ

// %operand: [[1, 2, 3], [4, 5, 6]]]
%result = "stablehlo.reshape"(%operand) : (tensor<2x3xi32>) -> tensor<3x2xi32>
// %result: [[1, 2], [3, 4], [5, 6]]

Ví dụ khác

đảo ngược

Ngữ nghĩa

Đảo ngược thứ tự các phần tử trong operand dọc theo dimensions được chỉ định và tạo ra một tensor result. Chính thức hơn, result[result_index] = operand[operand_index], trong đó:

  • operand_index[d] = dim(result, d) - result_index[d] - 1 nếu là d trong dimensions.
  • Nếu không, operand_index[d] = result_index[d].

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C3)
(I2) dimensions Hằng số tensor 1 chiều thuộc loại si64 (C2), (C3)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C3)

Giới hạn

  • (C1) type(operand) = type(result).
  • (C2) is_unique(dimensions).
  • (C3) 0 <= dimensions < rank(result).

Ví dụ

// %operand = [[1, 2], [3, 4], [5, 6]]
%result = "stablehlo.reverse"(%operand) {
  dimensions = dense<1> : tensor<1xi64>
} : (tensor<3x2xi32>) -> tensor<3x2xi32>
// %result: [[2, 1], [4, 3], [6, 5]]

Ví dụ khác

rng

Ngữ nghĩa

Tạo các số ngẫu nhiên bằng thuật toán rng_distribution và tạo ra tensor result của một hình dạng shape cho trước.

Nếu là rng_distribution = UNIFORM, thì các số ngẫu nhiên sẽ được tạo theo sự phân phối đồng nhất trong khoảng [a, b). Nếu là a >= b, hành vi sẽ không xác định.

Nếu là rng_distribution = NORMAL, thì các số ngẫu nhiên sẽ được tạo theo phân phối chuẩn với giá trị trung bình = a và độ lệch chuẩn = b. Nếu là b < 0, hành vi này chưa được xác định.

Cách chính xác về cách tạo số ngẫu nhiên là do triển khai xác định. Ví dụ: các trạng thái này có thể mang hoặc không mang tính quyết định, và có thể sử dụng hoặc không sử dụng trạng thái ẩn.

Trong các cuộc trò chuyện với nhiều bên liên quan, op này đã bị ngừng sử dụng một cách hiệu quả, vì vậy trong tương lai, chúng tôi dự định tìm hiểu cách xoá nó (#597).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) a Tensor 0 chiều của loại số nguyên, boolean hoặc dấu phẩy động (C1), (C2)
(I2) b Tensor 0 chiều của loại số nguyên, boolean hoặc dấu phẩy động (C1), (C2)
(I3) shape Hằng số tensor 1 chiều thuộc loại si64 (C3)
(I4) rng_distribution enum của UNIFORMNORMAL (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên, boolean hoặc dấu phẩy động (C1–C3)

Giới hạn

  • (C1) element_type(a) = element_type(b) = element_type(result).
  • (C2) Nếu là rng_distribution = NORMAL, thì is_float(a).
  • (C3) shape(result) = shape.

Ví dụ

// %a = 0
// %b = 2
// %shape = [3, 3]
%result = "stablehlo.rng"(%a, %b, %shape) {
  rng_distribution = #stablehlo<rng_distribution UNIFORM>
} : (tensor<i32>, tensor<i32>, tensor<2xi64>) -> tensor<3x3xi32>
// %result: [
//           [1, 0, 1],
//           [1, 1, 1],
//           [0, 0, 0]
//          ]

rng_bit_generator

Ngữ nghĩa

Trả về một output chứa các bit ngẫu nhiên đồng nhất và trạng thái đầu ra đã cập nhật output_state bằng cách sử dụng thuật toán trình tạo số giả ngẫu nhiên rng_algorithm khi có trạng thái ban đầu initial_state. Kết quả đầu ra được đảm bảo là hàm xác định của initial_state, nhưng không được đảm bảo là có tính xác định giữa các lần triển khai.

rng_algorithm có một trong những trạng thái sau đây:

  • DEFAULT: Thuật toán do việc triển khai xác định.
  • THREE_FRY: Biến thể do quy trình triển khai xác định của thuật toán Threefry.*
  • PHILOX: Biến thể do việc triển khai xác định của thuật toán Philox.*

* Xem: Salmon et al. SC 2011. Số ngẫu nhiên song song: dễ như 1, 2, 3.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) rng_algorithm enum của DEFAULT, THREE_FRYPHILOX (C2)
(I2) initial_state Tensor 1 chiều thuộc loại ui64 (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
output_state Tensor 1 chiều thuộc loại ui64 (C1)
output tensor của loại số nguyên hoặc dấu phẩy động

Giới hạn

  • (C1) type(initial_state) = type(output_state).
  • (C2) size(initial_state) được định nghĩa là:
    • do rng_algorithm = DEFAULT xác định.
    • 2 nếu rng_algorithm = THREE_FRY.
    • 2 hoặc 3 nếu là rng_algorithm = PHILOX.

Ví dụ

// %initial_state: [1, 2]
%output_state, %output = "stablehlo.rng_bit_generator"(%initial_state) {
  rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
} : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x2xui64>)
// %output_state: [1, 6]
// %output: [
//           [9236835810183407956, 16087790271692313299],
//           [18212823393184779219, 2658481902456610144]
//          ]

round_nearest_afz

Ngữ nghĩa

Thực hiện làm tròn phần tử theo số nguyên gần nhất, phá vỡ các mối quan hệ từ 0 trên tensor operand và tạo ra một tensor result. Triển khai thao tác roundToIntegralTiesToAway từ thông số kỹ thuật IEEE-754. Đối với các loại lượng tử, hãy thực hiện dequantize_op_quantize(round_nearest_afz, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_afz"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-3.0, 0.0, 1.0, 1.0, 3.0]

Ví dụ khác

round_nearest_even

Ngữ nghĩa

Thực hiện làm tròn theo phần tử hướng tới số nguyên gần nhất, phá vỡ các mối liên kết với số nguyên chẵn trên tensor operand và tạo ra một tensor result. Triển khai thao tác roundToIntegralTiesToEven từ quy cách IEEE-754. Đối với các loại lượng tử hoá, thực hiện dequantize_op_quantize(round_nearest_even, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor thuộc loại dấu phẩy động hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand = [-2.5, 0.4, 0.5, 0.6, 2.5]
%result = "stablehlo.round_nearest_even"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// %result: [-2.0, 0.0, 0.0, 1.0, 2.0]

Ví dụ khác

giải đấu

Ngữ nghĩa

Thực hiện thao tác căn bậc hai đối xứng 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: rSqrt từ IEEE-754.
  • Đối với số phức: căn bậc hai nghịch đảo của số phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(rsqrt, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [[1.0, 4.0], [9.0, 25.0]]
%result = "stablehlo.rsqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.5], [0.33333343, 0.2]]

Ví dụ khác

scatter

Ngữ nghĩa

Tạo tensor results bằng tensor inputs ngoại trừ việc một số lát do scatter_indices chỉ định được cập nhật bằng các giá trị updates bằng update_computation.

Sơ đồ sau đây cho thấy cách các phần tử trong updates... ánh xạ với các phần tử trong results... bằng một ví dụ cụ thể. Sơ đồ này chọn một vài chỉ mục updates... ví dụ và giải thích chi tiết những chỉ mục results... tương ứng.

Chính thức hơn, cho tất cả update_index trong index_space(updates[0]):

  • update_scatter_dims = [d for d in axes(updates[0]) and d not in update_window_dims].
  • update_scatter_index = update_index[update_scatter_dims...].
  • start_index được khai báo là:
    • scatter_indices[si0, ..., :, ..., siN], trong đó si là các phần tử riêng lẻ trong update_scatter_index: được chèn vào chỉ mục index_vector_dim, nếu index_vector_dim < rank(scatter_indices).
    • Nếu không, [scatter_indices[update_scatter_index]].
  • Đối với d_input trong axes(inputs[0]),
    • full_start_index[d_input] = start_index[d_start] nếu d_input = scatter_dims_to_operand_dims[d_start].
    • Nếu không, full_start_index[d_input] = 0.
  • update_window_index = update_index[update_window_dims...].
  • full_window_index = [wi0, ..., 0, ..., wiN], trong đó wi là các phần tử riêng lẻ trong update_window_index0 được chèn tại các chỉ mục từ inserted_window_dims.
  • result_index = full_start_index + full_window_index.

Do đó, results = exec(schedule, inputs), trong đó:

  • schedule là một hoán vị được xác định dựa trên quá trình triển khai của index_space(updates[0]).
  • exec([update_index, ...], results) = exec([...], updated_results) trong đó:
    • Nếu result_index nằm trong giới hạn của shape(results...)
    • updates_converted = to_destination_type( updates...[update_index], type(func_inputs(update_computation) [len(func_inputs(update_computation))//2:])... )
    • updated_values = update_computation(results...[result_index], updates_converted)
    • updated_results là bản sao của results với results...[result_index] được đặt thành updated_values....
    • Hoặc
    • updated_results = results.
  • exec([], results) = results.

Nếu indices_are_sortedtrue thì quá trình triển khai có thể giả định rằng scatter_indices được sắp xếp theo scatter_dims_to_operand_dims, nếu không thì hành vi sẽ không được xác định. Chính thức hơn, đối với tất cả i1 < i2 từ indices(result), full_start_index(i1) <= full_start_index(i2).

Nếu unique_indicestrue thì quá trình triển khai có thể giả định rằng tất cả chỉ mục result_index phân tán là duy nhất. Nếu unique_indicestrue nhưng các chỉ mục đang được phân tán không phải là duy nhất, thì hành vi sẽ không được xác định.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1), (C2), (C4-C6), (C10), (C13), (C15-C16)
(I2) scatter_indices tensor của kiểu số nguyên (C4), (C11), (C14)
(I3) updates số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C3-C6), (C8)
(I4) update_window_dims Hằng số tensor 1 chiều thuộc loại si64 (C2), (C4), (C7), (C8)
(I5) inserted_window_dims Hằng số tensor 1 chiều thuộc loại si64 (C2), (C4), (C9), (C10)
(I6) scatter_dims_to_operand_dims Hằng số tensor 1 chiều thuộc loại si64 (C11–C13)
(I7) index_vector_dim hằng số có kiểu si64 (C4), (C11), (C14)
(I8) indices_are_sorted hằng số có kiểu i1
(I9) unique_indices hằng số có kiểu i1
(I10) update_computation hàm (C15)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C15–C17)

Giới hạn

  • (C1) same(shape(inputs...)).
  • (C2) rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims).
  • (C3) same(shape(updates...)).
  • (C4) shape(updates[0]) = combine(update_scatter_dim_sizes, update_window_dim_sizes) trong đó:
    • update_scatter_dim_sizes = shape(scatter_indices) ngoại trừ việc kích thước kích thước của scatter_indices tương ứng với index_vector_dim sẽ không được đưa vào.
    • update_window_dim_sizes <= shape(inputs[0]) ngoại trừ việc kích thước kích thước trong inputs[0] tương ứng với inserted_window_dims sẽ không được đưa vào.
    • combine đặt update_scatter_dim_sizes ở các trục tương ứng với update_scatter_dimsupdate_window_dim_sizes tại các trục tương ứng với update_window_dims.
  • (C5) 0 < size(inputs) = size(updates) = N.
  • (C6) element_type(updates...) = element_type(inputs...).
  • (C7) is_unique(update_window_dims) and is_sorted(update_window_dims).
  • (C8) 0 <= update_window_dims < rank(updates[0]).
  • (C9) is_unique(inserted_window_dims) and is_sorted(update_window_dims).
  • (C10) 0 <= inserted_window_dims < rank(inputs[0]).
  • (C11) size(scatter_dims_to_operand_dims) = index_vector_dim < rank(scatter_indices) ? dim(scatter_indices, index_vector_dim) : 1.
  • (C12) is_unique(scatter_dims_to_operand_dims).
  • (C13) 0 <= scatter_dims_to_operand_dims < rank(inputs[0]).
  • (C14) 0 <= index_vector_dim <= rank(scatter_indices).
  • (C15) update_computation có loại (tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>), trong đó is_promotable(element_type(inputs[i]), Ei).
  • (C16) shape(inputs...) = shape(results...).
  • (C17) element_type(results[i]) = Ei cho tất cả i trong [0,N).

Ví dụ

// %input: [
//          [[1, 2], [3, 4], [5, 6], [7, 8]],
//          [[9, 10], [11, 12], [13, 14], [15, 16]],
//          [[17, 18], [19, 20], [21, 22], [23, 24]]
//         ]
// %scatter_indices: [[[0, 2], [1, 0], [2, 1]], [[0, 1], [1, 0], [0, 9]]]
// %update: [
//           [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]],
//           [[[1, 1], [1, 1]], [[1, 1], [1, 1]], [[1, 1], [1, 1]]]
//          ]
%result = "stablehlo.scatter"(%input, %scatter_indices, %update) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  scatter_dimension_numbers = #stablehlo.scatter<
    update_window_dims = [2, 3],
    inserted_window_dims = [0],
    scatter_dims_to_operand_dims = [1, 0],
    index_vector_dim = 2>,
  indices_are_sorted = false,
  unique_indices = false
} : (tensor<3x4x2xi64>, tensor<2x3x2xi64>, tensor<2x3x2x2xi64>) -> tensor<3x4x2xi64>
// %result: [
//           [[1, 2], [5, 6], [7, 8], [7, 8]],
//           [[10, 11], [12, 13], [14, 15], [16, 17]],
//           [[18, 19], [20, 21], [21, 22], [23, 24]]
//          ]

Ví dụ khác

chọn

Ngữ nghĩa

Tạo một tensor result, trong đó mỗi phần tử được chọn từ tensor on_true hoặc on_false dựa trên giá trị của phần tử pred tương ứng. Chính thức hơn là result[result_index] = pred_element ? on_true[result_index] : on_false[result_index], trong đó pred_element = rank(pred) = 0 ? pred[] : pred[result_index]. Đối với các loại lượng tử hoá, thực hiện dequantize_select_quantize(pred, on_true, on_false, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) pred tensor thuộc loại i1 (C1)
(I2) on_true Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C2)
(I3) on_false Tensor lượng tử hoá tensor hoặc mỗi tensor (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C2)

Giới hạn

  • (C1) rank(pred) = 0 or shape(pred) = shape(on_true).
  • (C2) baseline_type(on_true) = baseline_type(on_false) = baseline_type(result).

Ví dụ

// %pred: [[false, true], [true, false]]
// %on_true: [[1, 2], [3, 4]]
// %on_false: [[5, 6], [7, 8]]
%result = "stablehlo.select"(%pred, %on_true, %on_false) : (tensor<2x2xi1>, tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[5, 2], [3, 8]]

Ví dụ khác

select_and_scatter

Ngữ nghĩa

Phân tán các giá trị từ tensor source bằng cách sử dụng scatter dựa trên kết quả của reduce_window của tensor input bằng cách sử dụng select và tạo ra một tensor result.

Sơ đồ sau đây cho thấy cách tính toán các phần tử trong result từ operandsource, sử dụng một ví dụ cụ thể.

Chính thức hơn:

  • selected_values = reduce_window_without_init(...) với dữ liệu đầu vào sau:

    • `inputs = [toán hạng].
    • window_dimensions, window_stridespadding được sử dụng nguyên trạng.
    • base_dilations = windows_dilations = 1.
    • body được xác định là:
    def body(arg0: tensor<E>, arg1: tensor<E>) -> tensor<E>:
      return select(arg0, arg1) ? arg0 : arg1;
    

    trong đó E = element_type(operand)reduce_window_without_init hoạt động giống hệt như reduce_window, ngoại trừ việc schedule của reduce cơ bản (xem giảm) không bao gồm các giá trị khởi tạo. Hiện tại, bạn chưa chỉ định điều gì sẽ xảy ra nếu cửa sổ tương ứng không có giá trị (#731).

  • result[result_index] = reduce([source_values], [init_value], [0], scatter) trong đó:

    • source_values = [source[source_index] for source_index in source_indices].
    • selected_index(source_index) = operand_index nếu selected_values[source_index] có phần tử operand qua operand_index.
    • source_indices = [source_index for source_index in indices(source) if selected_index(source_index) = result_index].

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C4), (C6), (C8–C11)
(I2) source Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C2)
(I3) init_value Tensor lượng tử hoá tensor 0 chiều hoặc mỗi tensor (C3)
(I4) window_dimensions Hằng số tensor 1 chiều thuộc loại si64 (C2), (C4), (C5)
(I5) window_strides Hằng số tensor 1 chiều thuộc loại si64 (C2), (C6), (C7)
(I6) padding Hằng số tensor 2 chiều thuộc loại si64 (C2), (C8)
(I7) select hàm (C9)
(I8) scatter hàm (C10)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C11–C12)

Giới hạn

  • (C1) element_type(operand) = element_type(source).
  • (C2) shape(source) = num_windows trong đó:
    • padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1].
    • is_empty_window = padded_operand_shape = 0 || window_dimensions > padded_operand_shape.
    • num_windows = is_empty_window ? 0 : floor((padded_operand_shape - window_dimensions) / window_strides) + 1.
  • (C3) element_type(init_value) = element_type(operand).
  • (C4) size(window_dimensions) = rank(operand).
  • (C5) 0 < window_dimensions.
  • (C6) size(window_strides) = rank(operand).
  • (C7) 0 < window_strides.
  • (C8) shape(padding) = [rank(operand), 2].
  • (C9) select có loại (tensor<E>, tensor<E>) -> tensor<i1>, trong đó E = element_type(operand).
  • (C10) scatter có loại (tensor<E>, tensor<E>) -> tensor<E>, trong đó is_promotable(element_type(operand), E).
  • (C11) shape(operand) = shape(result).
  • (C12) element_type(result) = E.

Ví dụ

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

Ví dụ khác

gửi

Ngữ nghĩa

Gửi inputs đến kênh channel_id và tạo mã thông báo result.

Nếu is_host_transfertrue, thì thao tác này sẽ chuyển dữ liệu sang máy chủ lưu trữ. Nếu không, thiết bị sẽ chuyển dữ liệu sang một thiết bị khác. Điều này có nghĩa là việc triển khai được xác định. Cờ này sao chép thông tin được cung cấp trong channel_type, vì vậy, trong tương lai, chúng tôi dự định chỉ giữ lại một trong số đó (#666).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số variadic của tensor hoặc tensor lượng tử hoá
(I2) token token
(I3) channel_id hằng số có kiểu si64
(I4) channel_type enum của DEVICE_TO_DEVICEDEVICE_TO_HOST (C1)
(I5) is_host_transfer hằng số có kiểu i1 (C1)

Kết quả đầu ra

Tên Loại
result token

Giới hạn

  • (C1) channel_type được định nghĩa là:
    • DEVICE_TO_HOST nếu is_host_transfer = true,
    • Nếu không, DEVICE_TO_DEVICE.

Ví dụ

%result = "stablehlo.send"(%operand, %token) {
  channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
  is_host_transfer = true
} : (tensor<2x2xi64>, !stablehlo.token) -> !stablehlo.token

Ví dụ khác

shift_left

Ngữ nghĩa

Thực hiện thao tác dịch chuyển sang trái theo phần tử trên tensor lhs bằng số lượng bit rhs và tạo ra một tensor result.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của kiểu số nguyên (C1)
(I2) rhs tensor của kiểu số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên (C1)

Giới hạn

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

Ví dụ

// %lhs: [-1, 0, 1]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_left"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-2, 0, 8]

Ví dụ khác

shift_right_arithmetic

Ngữ nghĩa

Thực hiện thao tác dịch chuyển phải số học theo số học trên tensor lhs theo số lượng bit rhs và tạo ra một tensor result.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của kiểu số nguyên (C1)
(I2) rhs tensor của kiểu số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên (C1)

Giới hạn

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

Ví dụ

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_arithmetic"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [-1, 0, 1]

Ví dụ khác

shift_right_logical

Ngữ nghĩa

Thực hiện thao tác dịch chuyển phải logic theo các phần tử trên tensor lhs theo số lượng bit rhs và tạo ra một tensor result.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của kiểu số nguyên (C1)
(I2) rhs tensor của kiểu số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu số nguyên (C1)

Giới hạn

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

Ví dụ

// %lhs: [-1, 0, 8]
// %rhs: [1, 2, 3]
%result = "stablehlo.shift_right_logical"(%lhs, %rhs): (tensor<3xi64>, tensor<3xi64>) -> tensor<3xi64>
// %result: [9223372036854775807, 0, 1]

Ví dụ khác

biển báo

Ngữ nghĩa

Trả về dấu của phần tử operand và tạo ra tensor result. Chính thức hơn, đối với mỗi phần tử x, ngữ nghĩa có thể được biểu thị bằng cú pháp Python như sau:

def sign(x):
  if is_integer(x):
    if compare(x, 0, LT, SIGNED): return -1
    if compare(x, 0, EQ, SIGNED): return 0
    return 1
  elif is_float(x):
    if is_nan(x): return NaN
    if compare(x, -0.0, EQ, FLOAT): return -0.0
    if compare(x, +0.0, EQ, FLOAT): return +0.0
    if compare(x, 0.0, LT, FLOAT): return -1.0
    return 1.0
  elif is_complex(x):
    if is_nan(real(x)) or is_nan(imag(x)): return (NaN, NaN)
    if compare(x, (0.0, 0.0), EQ, FLOAT): return (0.0, 0.0)
    return divide(x, convert(abs(x), type(x)))

Đối với các loại lượng tử hoá, thực hiện dequantize_op_quantize(sign, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của số nguyên có dấu, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của số nguyên có dấu, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// operand: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]
%result = "stablehlo.sign"(%operand) : (tensor<5xf64>) -> tensor<5xf64>
// Logical values: +NaN, -1.0, -0.0, +0.0, 1.0
// %result: [0x7FFFFFFFFFFFFFFF, -1.0, -0.0, 0.0, 1.0]

Ví dụ khác

sin

Ngữ nghĩa

Thực hiện phép toán sin 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: sin từ IEEE-754.
  • Đối với số phức: sin phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(sine, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

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.sine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [0.0, -1.0]]

Ví dụ khác

lát cắt

Ngữ nghĩa

Trích xuất một lát cắt từ operand bằng cách sử dụng các chỉ mục khởi động tính toán tĩnh và tạo ra một tensor result. start_indices chứa các chỉ mục bắt đầu của phần cắt cho mỗi chiều, limit_indices chứa các chỉ mục kết thúc (dành riêng cho phần cắt của mỗi chiều) và strides chứa các chỉ mục kết thúc cho mỗi chiều.

Chính thức hơn là result[result_index] = operand[operand_index], trong đó operand_index = start_indices + result_index * strides.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor lượng tử hoá tensor hoặc mỗi tensor (C1–C3), (C5)
(I2) start_indices Hằng số tensor 1 chiều thuộc loại si64 (C2), (C3), (C5)
(I3) limit_indices Hằng số tensor 1 chiều thuộc loại si64 (C2), (C3), (C5)
(I4) strides Hằng số tensor 1 chiều thuộc loại si64 (C2), (C4)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor lượng tử hoá tensor hoặc mỗi tensor (C1), (C5)

Giới hạn

  • (C1) element_type(operand) = element_type(result).
  • (C2) size(start_indices) = size(limit_indices) = size(strides) = rank(operand).
  • (C3) 0 <= start_indices <= limit_indices <= shape(operand).
  • (C4) 0 < strides.
  • (C5) shape(result) = ceil((limit_indices - start_indices) / strides).

Ví dụ

// %operand: [
//            [0, 0, 0, 0],
//            [0, 0, 1, 1],
//            [0, 0, 1, 1]
//           ]
%result = "stablehlo.slice"(%operand) {
  start_indices = dense<[1, 2]> : tensor<2xi64>,
  limit_indices = dense<[3, 4]> : tensor<2xi64>,
  strides = dense<1> : tensor<2xi64>
} : (tensor<3x4xi64>) -> tensor<2x2xi64>
// % result: [
//            [1, 1],
//            [1, 1]
//           ]

Ví dụ khác

sắp xếp

Ngữ nghĩa

Sắp xếp các lát cắt 1 chiều của inputs dọc theo chiều dimension với nhau, theo comparator và tạo ra results.

Không giống như dữ liệu đầu vào tương tự trong các toán tử khác, dimension cho phép các giá trị âm với các ngữ nghĩa được mô tả bên dưới. Trong tương lai, điều này có thể không được phép vì lý do nhất quán (#1377).

Nếu is_stable là đúng, thì cách sắp xếp sẽ ổn định, tức là thứ tự tương đối của các phần tử mà trình so sánh coi là bằng nhau sẽ được giữ nguyên. Trong trường hợp có một giá trị đầu vào duy nhất, hai phần tử e1e2 được trình so sánh coi là bằng nhau khi và chỉ khi comparator(e1, e2) = comparator(e2, e1) = false. Hãy xem quy trình chính thức bên dưới để biết cách khái quát hoá cho nhiều dữ liệu đầu vào.

Chính thức hơn, cho tất cả result_index trong index_space(results[0]):

  • adjusted_dimension = dimension >= 0 ? dimension : rank(inputs[0]) + dimension.
  • result_slice = [ri0, ..., :, ..., riR-1], trong đó riN là các phần tử riêng lẻ trong result_index: được chèn vào adjusted_dimension.
  • inputs_together = (inputs[0]..., ..., inputs[N-1]...).
  • results_together[result_slice] = sort(inputs_together[result_slice], comparator_together).
  • trong đó sort sắp xếp một lát cắt 1 chiều theo thứ tự không giảm dần, với kỳ vọng comparator_together sẽ trả về true nếu đối số bên trái nhỏ hơn đối số thứ hai bên phải.
  • def comparator_together(lhs_together, rhs_together):
      args = []
      for (lhs_el, rhs_el) in zip(lhs_together, rhs_together):
        args.append(lhs_el)
        args.append(rhs_el)
      return comparator(*args)
    
  • (results[0]..., ..., results[N-1]...) = results_together.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) inputs số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C1–C5)
(I2) dimension hằng số có kiểu si64 (C4)
(I3) is_stable hằng số có kiểu i1
(I4) comparator hàm (C5)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên hoặc tensor lượng tử hoá mỗi tensor (C2), (C3)

Giới hạn

  • (C1) 0 < size(inputs).
  • (C2) type(inputs...) = type(results...).
  • (C3) same(shape(inputs...) + shape(results...)).
  • (C4) -R <= dimension < R, trong đó R = rank(inputs[0]).
  • (C5) comparator có loại (tensor<E1>, tensor<E1>, ..., tensor<EN-1>, tensor<EN-1>) -> tensor<i1>, trong đó Ei = element_type(inputs[i]).

Ví dụ

// %input0 = [[1, 2, 3], [3, 2, 1]]
// %input1 = [[3, 2, 1], [1, 2, 3]]
%result0, %result1 = "stablehlo.sort"(%input0, %input1) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>, %arg2: tensor<i64>, %arg3: tensor<i64>):
    %predicate = "stablehlo.compare"(%arg0, %arg1) {
      comparison_direction = #stablehlo<comparison_direction GT>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    "stablehlo.return"(%predicate) : (tensor<i1>) -> ()
}) {
  dimension = 0 : i64,
  is_stable = true
} : (tensor<2x3xi64>, tensor<2x3xi64>) -> (tensor<2x3xi64>, tensor<2x3xi64>)
// %result0 = [[3, 2, 3], [1, 2, 1]]
// %result1 = [[1, 2, 1], [3, 2, 3]]

Ví dụ khác

sqrt

Ngữ nghĩa

Thực hiện phép toán căn bậc hai của 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: squareRoot từ IEEE-754.
  • Đối với số phức: căn bậc hai của số phức.
  • Đối với các loại lượng tử hoá: dequantize_op_quantize(sqrt, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [[0.0, 1.0], [4.0, 9.0]]
%result = "stablehlo.sqrt"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[0.0, 1.0], [2.0, 3.0]]

Ví dụ khác

trừ

Ngữ nghĩa

Thực hiện phép trừ phần tử của hai tensor lhsrhs, đồ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ố nguyên: phép trừ số nguyên.
  • Đối với số thực: subtraction từ IEEE-754.
  • Đối với số phức: phép trừ phức.
  • Đối với các loại lượng tử hoá:
    • dequantize_op_quantize(subtract, lhs, rhs, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)
(I2) rhs tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của số nguyên, dấu phẩy động, loại phức hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

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

Ví dụ khác

tanh

Ngữ nghĩa

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

  • Đối với số thực: tanh từ IEEE-754.
  • Đối với số phức: tang hyperbol phức tạp.
  • Đối với các loại lượng tử hoá:
    • dequantize_op_quantize(tanh, operand, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

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

Ví dụ

// %operand: [-1.0, 0.0, 1.0]
%result = "stablehlo.tanh"(%operand) : (tensor<3xf32>) -> tensor<3xf32>
// %result: [-0.76159416, 0.0, 0.76159416]

Ví dụ khác

hoán vị

Ngữ nghĩa

Chuyển đổi kích thước của tensor operand bằng cách sử dụng permutation và tạo ra một tensor result. Chính thức hơn là result[result_index] = operand[operand_index], trong đó result_index[d] = operand_index[permutation[d]].

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand Tensor hoặc tensor lượng tử hoá (C1–C4)
(I2) permutation Hằng số tensor 1 chiều thuộc loại si64 (C2–C4)

Kết quả đầu ra

Tên Loại Giới hạn
result Tensor hoặc tensor lượng tử hoá (C1), (C3–C4)

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)quantization_dimension(result) có thể khác nhau, nếu không.
  • (C2) permutation là hoán vị của range(rank(operand)).
  • (C3) shape(result) = dim(operand, permutation...).
  • (C4) Nếu là is_per_axis_quantized(result), thì quantization_dimension(operand) = permutation(quantization_dimension(result)).

Ví dụ

// %operand: [
//            [[1,2], [3,4], [5,6]],
//            [[7,8], [9,10], [11,12]]
//           ]
%result = "stablehlo.transpose"(%operand) {
  permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
// %result: [
//           [[1,7], [3,9], [5,11]],
//           [[2,8], [4,10], [6,12]]
//          ]

Ví dụ khác

triangular_solve

Ngữ nghĩa

Giải các lô hệ thống phương trình tuyến tính có ma trận hệ số tam giác trên hoặc dưới.

Nói một cách chính thức hơn, với ab, result[i0, ..., iR-3, :, :] là giải pháp cho op(a[i0, ..., iR-3, :, :]) * x = b[i0, ..., iR-3, :, :] khi left_sidetrue hoặc x * op(a[i0, ..., iR-3, :, :]) = b[i0, ..., iR-3, :, :] khi left_sidefalse, giải quyết biến x trong đó op(a) được xác định bởi transpose_a, có thể là một trong những điều sau:

  • NO_TRANSPOSE: Thực hiện thao tác bằng cách sử dụng a nguyên trạng.
  • TRANSPOSE: Thực hiện thao tác khi hoán vị của a.
  • ADJOINT: Thực hiện thao tác trên hoán vị liên hợp của a.

Dữ liệu đầu vào chỉ được đọc từ tam giác dưới của a, nếu lowertrue hoặc tam giác trên của a, nếu không thì. Dữ liệu đầu ra được trả về trong cùng một tam giác; các giá trị trong tam giác khác được xác định theo phương thức triển khai.

Nếu unit_diagonal là đúng, thì quá trình triển khai có thể giả định rằng các phần tử đường chéo của a bằng 1, nếu không hành vi sẽ không xác định được.

Đối với các loại lượng tử hoá, thực hiện dequantize_op_quantize(lambda x, y: triangular_solve(x, y, left_side, lower, unit_diagonal, transpose_a), a, b, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) a tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1–C3)
(I2) b tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1–C4)
(I3) left_side hằng số có kiểu i1 (C3)
(I4) lower hằng số có kiểu i1
(I5) unit_diagonal hằng số có kiểu i1
(I6) transpose_a enum của NO_TRANSPOSE, TRANSPOSEADJOINT

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động hoặc phức tạp hoặc tensor lượng tử hoá mỗi tensor (C1)

Giới hạn

  • (C1) baseline_element_type(a) = baseline_element_type(b).
  • (C2) 2 <= rank(a) = rank(b) = R.
  • (C3) Mối quan hệ giữa shape(a)shape(b) được định nghĩa như sau:
    • shape(a)[:-3] = shape(b)[:-3].
    • dim(a, -2) = dim(a, -1) = dim(b, left_side ? -2 : -1).
  • (C4) baseline_type(b) = baseline_type(result).

Ví dụ

// %a = [
//       [1.0, 0.0, 0.0],
//       [2.0, 4.0, 0.0],
//       [3.0, 5.0, 6.0]
//      ]
// %b = [
//       [2.0, 0.0, 0.0],
//       [4.0, 8.0, 0.0],
//       [6.0, 10.0, 12.0]
//      ]
%result = "stablehlo.triangular_solve"(%a, %b) {
  left_side = true,
  lower = true,
  unit_diagonal = false,
  transpose_a = #stablehlo<transpose NO_TRANSPOSE>
} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<3x3xf32>
// %result: [
//           [2.0, 0.0, 0.0],
//           [0.0, 2.0, 0.0],
//           [0.0, 0.0, 2.0]
//          ]

tuple

Ngữ nghĩa

Tạo một bộ dữ liệu result từ các giá trị val.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) val số lượng giá trị biến thiên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tuple (C1)

Giới hạn

  • (C1) result có loại tuple<E0, ..., EN-1>, trong đó Ei = type(val[i]).

Ví dụ

// %val0: [1.0, 2.0]
// %val1: (3)
%result = "stablehlo.tuple"(%val0, %val1) : (tensor<2xf32>, tuple<tensor<i32>>) -> tuple<tensor<2xf32>, tuple<tensor<i32>>>
// %result: ([1.0, 2.0], (3))

Ví dụ khác

uniform_dequantize

Ngữ nghĩa

Thực hiện việc chuyển đổi tensor lượng tử hoá operand thành tensor dấu phẩy động result theo các tham số lượng tử hoá do loại operand xác định.

Chính thức hơn là result = dequantize(operand).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor lượng tử hoá (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của loại dấu phẩy động (C1), (C2)

Giới hạn

  • (C1) shape(operand) = shape(result).
  • (C2) element_type(result) = expressed_type(operand).

Ví dụ

// %operand: [10, 10]
%result = "stablehlo.uniform_dequantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2xf32>
// %result: [4.0, 15.0]

uniform_quantize

Ngữ nghĩa

Thực hiện việc chuyển đổi tensor dấu phẩy động hoặc tensor lượng tử hoá operand thành tensor lượng tử hoá result theo các tham số lượng tử hoá do loại result xác định theo các tham số lượng tử hoá.

Chính thức hơn,

  • Nếu is_float(operand):
    • result = quantize(operand, type(result)).
  • Nếu is_quantized(operand):
    • float_result = dequantize(operand).
    • result = quantize(float_result, type(result)).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand tensor của loại dấu phẩy động hoặc lượng tử hoá (C1), (C2)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor lượng tử hoá (C1), (C2)

Giới hạn

  • (C1) shape(operand) = shape(result).
  • (C2) expressed_type(result) = is_float(operand) ? element_type(operand) : expressed_type(operand).

Ví dụ

// %operand: [4.0, 15.0]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>
// %result: [10, 10]

// %operand: [10, 10]
%result = "stablehlo.uniform_quantize"(%operand) : (tensor<2x!quant.uniform<i8:f32:0, {0.1:-30,0.5:-20}>>) -> tensor<2x!quant.uniform<i8:f32:0, {0.1:-20,0.2:-30}>>
// %result: [20, 45]

trong khi

Ngữ nghĩa

Tạo đầu ra từ việc thực thi hàm body từ 0 lần trở lên trong khi hàm cond xuất ra true. Chính thức hơn, bạn có thể biểu thị ngữ nghĩa bằng cú pháp Python như sau:

internal_state = operand
while cond(*internal_state):
  internal_state = body(*internal_state)
results = internal_state

Hành vi của vòng lặp vô hạn là TBD (#383).

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) operand số lượng tensor biến thiên, tensor lượng tử hoá hoặc mã thông báo (C1–C3)
(I2) cond hàm (C1)
(I3) body hàm (C2)

Kết quả đầu ra

Tên Loại Giới hạn
results số lượng tensor biến thiên, tensor lượng tử hoá hoặc mã thông báo (C3)

Giới hạn

  • (C1) cond có loại (T0, ..., TN-1) -> tensor<i1>, trong đó Ti = type(operand[i]).
  • (C2) body có loại (T0, ..., TN-1) -> (T0, ..., TN-1), trong đó Ti = type(operand[i]).
  • (C3) type(results...) = type(operand...).

Ví dụ

// %init_i: 1
// %init_sum: 0
// %one: 1
// %ten: 10
%results0, %results1 = "stablehlo.while"(%init_i, %init_sum) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %cond = "stablehlo.compare"(%arg0, %ten) {
      comparison_direction = #stablehlo<comparison_direction LT>
    } : (tensor<i64>, tensor<i64>) -> tensor<i1>
    stablehlo.return %cond : tensor<i1>
  }, {
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %new_sum = stablehlo.add %arg1, %one : tensor<i64>
    %new_i = stablehlo.add %arg0, %one : tensor<i64>
    stablehlo.return %new_i, %new_sum : tensor<i64>, tensor<i64>
}) : (tensor<i64>, tensor<i64>) -> (tensor<i64>, tensor<i64>)
// %results0: 10
// %results1: 10

Ví dụ khác

xor

Ngữ nghĩa

Thực hiện XOR theo phần tử của hai tensor lhsrhs, đồ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 XOR.
  • Đối với số nguyên: bitwise XOR.

Thông tin đầu vào

Nhãn Tên Loại Giới hạn
(I1) lhs tensor của kiểu boolean hoặc số nguyên (C1)
(I2) rhs tensor của kiểu boolean hoặc số nguyên (C1)

Kết quả đầu ra

Tên Loại Giới hạn
result tensor của kiểu boolean hoặc số nguyên (C1)

Giới hạn

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

Ví dụ

// Bitwise operation with with integer tensors
// %lhs: [[1, 2], [3, 4]]
// %rhs: [[5, 6], [7, 8]]
%result = "stablehlo.xor"(%lhs, %rhs) : (tensor<2x2xi32>, tensor<2x2xi32>) -> tensor<2x2xi32>
// %result: [[4, 4], [4, 12]]

// Logical operation with with boolean tensors
// %lhs: [[false, false], [true, true]]
// %rhs: [[false, true], [false, true]]
%result = "stablehlo.xor"(%lhs, %rhs) : (tensor<2x2xi1>, tensor<2x2xi1>) -> tensor<2x2xi1>
// %result: [[false, true], [true, false]]

Thực thi

Thực thi tuần tự

Chương trình StableHLO được thực thi bằng cách cung cấp các giá trị đầu vào cho hàm main và các giá trị đầu ra tính toán. Giá trị đầu ra của một hàm được tính toán bằng cách thực thi biểu đồ hoạt động đã bị can thiệp vào hệ thống trong hoạt động return tương ứng.

Thứ tự thực thi được xác định theo phương thức triển khai, miễn là thứ tự đó phù hợp với luồng dữ liệu, tức là liệu hoạt động có được thực thi trước khi sử dụng hay không. Trong StableHLO, tất cả các hoạt động có hiệu ứng phụ đều sử dụng một mã thông báo và tạo ra một mã thông báo (nhiều mã thông báo có thể được ghép thành một mã thông qua after_all), vì vậy, thứ tự thực thi của hiệu ứng phụ cũng được điều chỉnh cho phù hợp với luồng dữ liệu. Các lệnh thực thi có thể có của chương trình ví dụ ở trên là %0%1%2%3%4return hoặc %3%0%1%2%4return.

Nói một cách chính thức hơn, quy trình StableHLO là sự kết hợp giữa: 1) chương trình StableHLO, 2) trạng thái hoạt động (chưa được thực thi, đã được thực thi) và 3) các giá trị trung gian mà quy trình đang xử lý. Quá trình này bắt đầu với các giá trị đầu vào cho hàm main, diễn ra thông qua biểu đồ hoạt động cập nhật trạng thái hoạt động và các giá trị trung gian, sau đó kết thúc với các giá trị đầu ra. Quy trình chính thức tiếp theo sẽ được TBD (#484) xác định.

Thực thi song song

Các chương trình StableHLO có thể được thực thi song song, được sắp xếp thành một lưới quy trình 2D của num_replicas bởi num_partitions (cả hai đều có loại ui32).

Trong lưới quy trình StableHLO, num_replicas * num_partitions các quy trình StableHLO đang được thực thi cùng lúc. Mỗi quy trình có một process_id = (replica_id, partition_id) duy nhất, trong đó replica_id trong replica_ids = range(num_replicas)partition_id trong partition_ids = range(num_partitions) đều có loại ui32.

Kích thước của lưới quy trình được xác định tĩnh cho mọi chương trình (trong tương lai, chúng tôi dự định đưa kích thước này vào một phần rõ ràng của các chương trình StableHLO #650) và vị trí trong lưới quy trình được xác định tĩnh cho mọi quy trình. Mỗi quy trình đều có quyền truy cập vào vị trí của nó trong lưới quy trình thông qua hoạt động replica_idpartition_id.

Trong lưới xử lý, tất cả các chương trình đều có thể giống nhau (theo kiểu "Một chương trình, Nhiều dữ liệu"), tất cả đều có thể khác nhau (trong kiểu "Nhiều chương trình, Nhiều dữ liệu") hoặc một kiểu nào đó ở giữa. Trong tương lai, chúng tôi dự định ra mắt tính năng hỗ trợ các thành ngữ khác để xác định các chương trình StableHLO song song, bao gồm cả GSPMD (#619).

Trong lưới quy trình, các quy trình hầu như độc lập với nhau – chúng có trạng thái hoạt động riêng biệt, giá trị đầu vào/trung gian/đầu ra riêng biệt và hầu hết các hoạt động được thực thi riêng biệt giữa các quy trình, ngoại trừ một số ít hoạt động chung được mô tả dưới đây.

Do quá trình thực thi hầu hết các hoạt động chỉ sử dụng các giá trị từ cùng một quy trình nên thường rõ ràng việc tham chiếu các giá trị này theo tên của chúng. Tuy nhiên, khi mô tả ngữ nghĩa của các hoạt động tập hợp, điều đó là không đủ, và điều đó làm phát sinh ký hiệu name@process_id để tham chiếu đến giá trị name trong một quy trình cụ thể. (Từ đó, bạn có thể xem name không đủ điều kiện dưới dạng viết tắt của name@(replica_id(), partition_id())).

Thứ tự thực thi giữa các quy trình là do quy trình triển khai xác định, ngoại trừ quá trình đồng bộ hoá do hoạt động giao tiếp điểm-điểm và hoạt động tập thể tạo ra như mô tả dưới đây.

Giao tiếp giữa các điểm

Các quy trình StableHLO có thể giao tiếp với nhau thông qua các kênh StableHLO. Kênh được biểu thị bằng một mã nhận dạng dương của loại si64. Thông qua nhiều hoạt động, bạn có thể gửi giá trị đến các kênh và nhận giá trị đó từ các kênh.

Cung cấp thêm thông tin, ví dụ: nguồn gốc của các mã nhận dạng kênh này, cách các chương trình quy trình nhận biết được chúng và loại hình đồng bộ hoá nào được giới thiệu sẽ là TBD (#484).

Liên lạc phát trực tuyến

Mọi quy trình StableHLO đều có quyền truy cập vào hai giao diện truyền trực tuyến:

  • Nguồn cấp dữ liệu có thể đọc được.
  • Nguồn cấp dữ liệu ngoài có thể được ghi vào.

Không giống như kênh được dùng để giao tiếp giữa các quy trình và do đó có các quy trình ở cả hai đầu, nguồn cấp dữ liệu và nguồn cấp dữ liệu ngoài có phương thức triển khai đầu khác được xác định.

Cung cấp thêm thông tin chính thức, ví dụ: cách hoạt động giao tiếp truyền trực tuyến ảnh hưởng đến thứ tự thực thi và loại đồng bộ hoá được giới thiệu sẽ là TBD (#484).

Hoạt động tập thể

Có 6 hoạt động tổng hợp trong StableHLO: all_gather, all_reduce, all_to_all, collective_broadcast, collective_permutereduce_scatter. Tất cả hoạt động này chia các quy trình trong lưới quy trình StableHLO thành các nhóm quy trình StableHLO và thực thi phép tính toán chung trong từng nhóm quy trình, độc lập với các nhóm quy trình khác.

Trong mỗi nhóm quy trình, hoạt động tập thể có thể tạo ra một rào cản đồng bộ hoá. Cung cấp thêm thông tin, ví dụ: giải thích cụ thể hơn về thời điểm chính xác của quá trình đồng bộ hoá này, cách thức chính xác các quy trình vượt qua rào cản này và điều gì xảy ra nếu chúng không vượt qua rào cản này sẽ là TBD (#484).

Nếu nhóm quy trình có liên quan đến hoạt động giao tiếp giữa các phân vùng, tức là có các quy trình trong nhóm quy trình có mã phân vùng khác nhau, thì việc thực thi hoạt động chung cần có một kênh và hoạt động tập thể phải cung cấp channel_id dương thuộc loại si64. Giao tiếp giữa các bản sao không cần kênh.

Các phép tính do hoạt động tập thể thực hiện là dành riêng cho từng hoạt động riêng lẻ và được mô tả trong các phần hoạt động riêng lẻ ở trên. Tuy nhiên, các chiến lược mà lưới quy trình được chia thành các nhóm quy trình được chia sẻ giữa các hoạt động này và được mô tả trong phần này. Chính thức hơn, StableHLO hỗ trợ 4 chiến lược sau.

cross_replica

Chỉ giao tiếp nhiều bản sao mới xảy ra trong từng nhóm quy trình. Chiến lược này sẽ sử dụng replica_groups – một danh sách các mã nhận dạng bản sao – và tính toán tích Descartes của replica_groups theo partition_ids. replica_groups phải có các phần tử duy nhất và bao gồm tất cả replica_ids. Chính thức hơn, sử dụng cú pháp Python:

def cross_replica(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
  for replica_group in replica_groups:
    for partition_id in partition_ids:
      process_group = []
      for replica_id in replica_group:
        process_group.append((replica_id, partition_id))
      yield process_group

Ví dụ: đối với replica_groups = [[0, 1], [2, 3]]num_partitions = 2, cross_replica sẽ tạo ra [[(0, 0), (1, 0)], [(0, 1), (1, 1)], [(2, 0), (3, 0)], [(2, 1), (3, 1)]].

cross_partition

Chỉ có hoạt động giao tiếp giữa các phân vùng diễn ra trong mỗi nhóm quy trình. Chiến lược này sẽ lấy partition_groups – một danh sách danh sách mã nhận dạng phân vùng – và tính toán tích Descartes của partition_groups theo replica_ids. partition_groups phải có các phần tử riêng biệt và bao gồm tất cả partition_ids. Chính thức hơn là sử dụng cú pháp Python:

def cross_partition(partition_groups: List[List[PartitionId]]) -> List[List[ProcessId]]:
  for partition_group in partition_groups:
    for replica_id in replica_ids:
      process_group = []
      for partition_id in partition_group:
        process_group.append((replica_id, partition_id))
      yield process_group

Ví dụ: đối với partition_groups = [[0, 1]]num_replicas = 4, cross_partition sẽ tạo ra [[(0, 0), (0, 1)], [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)]].

cross_replica_and_partition

Cả hoạt động giao tiếp nhiều bản sao và phân vùng đều có thể xảy ra trong từng nhóm quy trình. Chiến lược này sẽ sử dụng replica_groups – một danh sách các mã nhận dạng bản sao – và tính toán các tích Descartes của mỗi replica_group bằng partition_ids. replica_groups phải có các phần tử riêng biệt và bao gồm tất cả replica_ids. Chính thức hơn là sử dụng cú pháp Python:

def cross_replica_and_partition(replica_groups: List[List[ReplicaId]]) -> List[List[ProcessId]]:
  for replica_group in replica_groups:
    process_group = []
    for partition_id in partition_ids:
      for replica_id in replica_group:
        process_group.append((replica_id, partition_id))
    yield process_group

Ví dụ: đối với replica_groups = [[0, 1], [2, 3]]num_partitions = 2, cross_replica_and_partition sẽ tạo ra [[(0, 0), (1, 0), (0, 1), (1, 1)], [(2, 0), (3, 0), (2, 1), (3, 1)]].

flattened_ids

Chiến lược này sẽ lấy flattened_id_groups – một danh sách danh sách mã quy trình "đã làm phẳng" ở dạng replica_id * num_partitions + partition_id – và chuyển các mã này thành mã quy trình. flattened_id_groups phải có các phần tử riêng biệt và bao gồm mọi process_ids. Chính thức hơn là sử dụng cú pháp Python:

def flattened_ids(flattened_id_groups: List[List[ui32]]) -> List[List[ProcessId]]:
  for flattened_id_group in flattened_id_groups:
    process_group = []
    for flattened_id in flattened_id_group:
      replica_id = flattened_id // num_partitions
      partition_id = flattened_id % num_partitions
      process_group.append((replica_id, partition_id))
    yield process_group

Ví dụ: đối với flattened_id_groups = [[0, 1, 2, 3], [4, 5, 6, 7]], num_replicas = 4num_partitions = 2, flattened_ids sẽ tạo ra [[(0, 0), (0, 1), (1, 0), (1, 1)], [(2, 0), (2, 1), (3, 0), (3, 1)]].

Độ chính xác

Hiện tại, StableHLO không đảm bảo về độ chính xác bằng số, nhưng điều này có thể thay đổi trong tương lai (#1156).

Lỗi

Các chương trình StableHLO được xác thực thông qua một tập hợp các quy tắc ràng buộc mở rộng cho từng hoạt động riêng lẻ, giúp loại bỏ nhiều lớp lỗi trước thời gian chạy. Tuy nhiên, các điều kiện lỗi vẫn có thể xảy ra, chẳng hạn như thông qua tràn số nguyên, quyền truy cập ngoài giới hạn, v.v. Trừ phi được gọi rõ ràng, tất cả các lỗi này đều dẫn đến hành vi do phương thức triển khai xác định, nhưng điều này có thể thay đổi trong tương lai (#1157).

Theo ngoại lệ với quy tắc này, các ngoại lệ dấu phẩy động trong các chương trình StableHLO có hành vi được xác định rõ ràng. Các thao tác dẫn đến các ngoại lệ do tiêu chuẩn IEEE-754 xác định (phép toán không hợp lệ, chia cho 0, tràn, luồng dưới hoặc ngoại lệ không chính xác) sẽ tạo ra kết quả mặc định (như đã xác định trong tiêu chuẩn) và tiếp tục thực thi mà không nâng cờ trạng thái tương ứng; tương tự như cách xử lý ngoại lệ raiseNoFlag từ tiêu chuẩn. Các trường hợp ngoại lệ đối với các phép toán không chuẩn (ví dụ: số học phức tạp và một số hàm siêu việt) sẽ do phương thức triển khai xác định.

Ký hiệu

Để mô tả cú pháp, tài liệu này đang sử dụng phiên bản ISO đã sửa đổi của cú pháp EBNF (ISO/IEC 14977:1996, Wikipedia), với hai sửa đổi: 1) quy tắc được xác định bằng ::= thay vì =,

2) phép nối được biểu thị bằng cách sử dụng phép nối thay vì ,.

Để mô tả ngữ nghĩa (tức là trong các phần "Loại", "Hằng số" và "Hoạt động"), chúng tôi đang sử dụng các công thức dựa trên cú pháp Python được mở rộng với khả năng hỗ trợ biểu thị ngắn gọn các thao tác mảng như mô tả dưới đây. Cách này phù hợp với các đoạn mã nhỏ, nhưng trong một số ít trường hợp khi cần đoạn mã lớn hơn, chúng tôi sử dụng cú pháp Python vanilla luôn được giới thiệu một cách rõ ràng.

Công thức

Hãy cùng khám phá cách hoạt động của công thức dựa trên một ví dụ từ quy cách dot_general. Một trong các quy tắc ràng buộc của thao tác này sẽ như sau: dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).

Tên dùng trong công thức này đến từ 2 nguồn: 1) hàm toàn cục (Global function) tức là dim, 2) định nghĩa thành phần của phần tử chương trình tương ứng, cụ thể là dữ liệu đầu vào lhs, lhs_batching_dimensions, rhsrhs_batching_dimensions được xác định trong phần "Đầu vào" của dot_general.

Như đã đề cập ở trên, cú pháp của công thức này dựa trên Python với một số phần mở rộng hướng tới tính súc tích. Để hiểu công thức này, hãy chuyển đổi thành cú pháp Python vanilla.

A) Trong các công thức này, chúng tôi đang sử dụng = để biểu thị tính đẳng thức, vì vậy, bước đầu tiên để có được cú pháp Python là thay = bằng ==, như sau: dim(lhs, lhs_batching_dimensions...) == dim(rhs, rhs_batching_dimensions...).

B) Ngoài ra, các công thức này hỗ trợ dấu ba chấm (...) để biến biểu thức vô hướng thành biểu thức tensor. Tóm lại, f(xs...) có nghĩa là "đối với mỗi x vô hướng trong tensor xs, hãy tính toán f(x) vô hướng rồi trả về tất cả các kết quả vô hướng này cùng nhau dưới dạng kết quả tensor". Trong cú pháp Python vanilla, công thức ví dụ của chúng ta sẽ chuyển thành: [dim(lhs, dim1) for dim1 in lhs_batching_dimensions] == [dim(rhs, dim2) for dim2 in rhs_batching_dimensions].

Nhờ có dấu ba chấm, bạn thường có thể tránh làm việc ở cấp độ vô hướng riêng lẻ. Tuy nhiên, trong một số trường hợp phức tạp, bạn có thể sử dụng cú pháp bán chính thức cấp thấp hơn như trong công thức start_indices[bi0, ..., :, ..., biN] từ quy cách gather. Để đảm bảo tính súc tích, chúng tôi không đưa ra hình thức chính xác để dịch cú pháp như vậy sang Python vanilla, hy vọng rằng theo từng trường hợp, cú pháp này vẫn có thể hiểu được một cách trực quan. Vui lòng cho chúng tôi biết nếu một số công thức cụ thể trông không rõ ràng và chúng tôi sẽ cố gắng cải thiện các công thức đó.

Ngoài ra, bạn sẽ nhận thấy các công thức sử dụng dấu ba chấm để mở rộng mọi loại danh sách, bao gồm cả tensor, danh sách tensor (ví dụ: có thể phát sinh từ một số tensor biến thiên), v.v. Đây là một khía cạnh khác mà chúng tôi không đưa ra hình thức chính xác (ví dụ: các danh sách thậm chí không phải là một phần của hệ thống kiểu StableHLO) và dựa vào khả năng hiểu.

C) Phương tiện ký hiệu đáng chú ý cuối cùng mà chúng tôi sử dụng đang truyền phát ngầm. Mặc dù phần tử điều khiển StableHLO không hỗ trợ tính năng thông báo truyền phát ngầm, nhưng các công thức này cũng hỗ trợ nhằm đảm bảo tính súc tích. Tóm lại, nếu đại lượng vô hướng được sử dụng trong ngữ cảnh dự kiến tensor, thì đại lượng vô hướng sẽ được truyền đến hình dạng dự kiến.

Để tiếp tục ví dụ về dot_general, dưới đây là một quy tắc ràng buộc khác: 0 <= lhs_batching_dimensions < rank(lhs). Theo định nghĩa trong quy cách dot_general, lhs_batching_dimensions là một tensor, tuy nhiên, cả 0rank(lhs) đều là đại lượng vô hướng. Sau khi chúng ta áp dụng tính năng thông báo truyền phát ngầm, công thức sẽ trở thành [0, ..., 0] <= lhs_batching_dimensions < [rank(lhs), ..., rank(lhs)].

Khi áp dụng cho một toán tử dot_general cụ thể, công thức này sẽ đánh giá một tensor của boolean. Khi công thức được dùng làm điều kiện ràng buộc, quy tắc ràng buộc sẽ được giữ nếu công thức đánh giá là true hoặc là một tensor chỉ có phần tử true.

Tên

Trong các công thức, phạm vi từ vựng bao gồm: 1) hàm toàn cục, 2) định nghĩa thành phần,

3) định nghĩa địa phương. Danh sách các hàm chung được cung cấp dưới đây. Danh sách định nghĩa phần tử phụ thuộc vào phần tử chương trình mà ký hiệu đó được áp dụng:

  • Đối với các toán tử, định nghĩa thành phần bao gồm tên được giới thiệu trong các phần "Đầu vào" và "Đầu ra".
  • Đối với mọi trường hợp khác, định nghĩa thành phần bao gồm các phần cấu trúc của phần tử chương trình, được đặt tên theo các đầu cuối tương ứng của EBNF. Trong hầu hết trường hợp, tên của các phần cấu trúc này được lấy bằng cách chuyển đổi tên của các điểm cuối không phải đầu cuối sang kiểu chữ thường (ví dụ: IntegerLiteral => integer_literal), nhưng đôi khi tên được viết tắt trong quá trình này (ví dụ: QuantizationStorageType => storage_type), trong trường hợp đó, tên được giới thiệu rõ ràng tương tự như các phần "Đầu vào" / "Đầu ra" trong quy cách hoạt động.
  • Ngoài ra, các định nghĩa thành phần luôn bao gồm self để tham chiếu đến phần tử tương ứng của chương trình.

Giá trị

Khi được đánh giá, công thức sẽ hoạt động với các loại giá trị sau: 1) Value (giá trị thực tế, ví dụ: dense<[[1, 2], [3, 4]]> : tensor<2x2xi32>; luôn biết loại của chúng), 2) Placeholder (các giá trị tương lai, ví dụ: lhs, rhs hoặc result; giá trị thực tế của chúng chưa được xác định, chỉ có các loại mới được biết), 3) Type (các loại như đã xác định trong phần "Loại"), 4) Function (các hàm chung như đã định nghĩa trong phần "Hàm").

Tuỳ thuộc vào ngữ cảnh, tên có thể đề cập đến các giá trị khác nhau. Cụ thể hơn, phần "Ngữ nghĩa" cho hoạt động (và mục tương đương cho các phần tử khác của chương trình) xác định logic thời gian chạy, vì vậy, tất cả dữ liệu đầu vào đều có sẵn dưới dạng Value. Ngược lại, phần "Quy tắc ràng buộc" đối với hoạt động (và các hoạt động tương đương) xác định logic "thời gian biên dịch", tức là một hoạt động nào đó thường được thực thi trước thời gian chạy, vì vậy, chỉ có dữ liệu đầu vào không đổi là có sẵn dưới dạng Value và các dữ liệu đầu vào khác chỉ có sẵn dưới dạng Placeholder.

Tên Trong "Ngữ nghĩa" Trong "Constraints" (Hạn chế)
Hàm toàn cục Function Function
Đầu vào không đổi Value Value
Dữ liệu đầu vào không cố định Value Placeholder
Kết quả đầu ra Value Placeholder
Định nghĩa địa phương Tuỳ thuộc vào định nghĩa Tuỳ thuộc vào định nghĩa

Hãy xem xét một thao tác transpose mẫu:

%result = "stablehlo.transpose"(%operand) {
  permutation = dense<[2, 1, 0]> : tensor<3xi64>
} : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>

Đối với toán tử này, permutation là một hằng số nên có sẵn dưới dạng Value trong cả ngữ nghĩa và quy tắc ràng buộc. Ngược lại, operandresult có sẵn dưới dạng Value về mặt ngữ nghĩa nhưng chỉ ở dạng Placeholder trong các quy tắc ràng buộc.

Hàm

Cấu trúc của các loại

Không có hàm nào có thể dùng để tạo loại. Thay vào đó, chúng tôi trực tiếp sử dụng cú pháp loại vì cú pháp này thường ngắn gọn hơn. Ví dụ: (tensor<E>, tensor<E>) -> (tensor<E>) thay vì function_type( [tensor_type([], E), tensor_type([], E)], [tensor_type([], E)]).

Hàm trên loại

  • element_type được xác định tương ứng trên các loại tensor, các loại tensor lượng tử hoá và các dữ liệu trả về, phần TensorElementType hoặc QuantizedTensorElementType của TensorType hoặc QuantizedTensorType tương ứng.
def element_type(x: Value | Placeholder | Type):
 if type(x) == TensorType:
    return tensor_element_type(x)
  if type(x) == QuantizedTensorType:
    return quantized_tensor_element_type(x)
  if type(x) is not Type:
    return element_type(type(x))
  • is_per_axis_quantized(x: Value | Placeholder | Type) -> Value là lối tắt cho is_quantized(x) and quantization_dimension(x) is not None.

  • is_per_tensor_quantized(x: Value | Placeholder | Type) -> Value là phím tắt cho is_quantized(x) and quantization_dimension(x) is None.

  • is_promotable(x: Type, y: Type) -> bool kiểm tra xem liệu loại x có thể được quảng bá thành loại y hay không. Khi xyQuantizedTensorElementType, chương trình khuyến mãi chỉ áp dụng cho storage_type. Phiên bản chương trình khuyến mãi cụ thể này hiện được dùng trong ngữ cảnh tính toán rút gọn (tham khảo RFC để biết thêm chi tiết).

def is_promotable(x: Type, y: Type) -> Value:
  is_same_type = (is_bool(x) and is_bool(y)) or
    (is_integer(x) and is_integer(y)) or (is_float(x) and is_float(y)) or
    (is_complex(x) and is_complex(y)) or
    (is_quantized(x) and is_quantized(y) and expressed_type(x) = expressed_type(y))

  if is_same_type == False:
    return False

  if is_integer(x) or is_float(x):
    return bitwidth(x) <= bitwidth(y)

  if is_complex(x):
    return bitwidth(element_type(x)) <= bitwidth(element_type(y))

  if is_quantized(x):
    return bitwidth(storage_type(x)) <= bitwidth(storage_type(y))

  return false
  • is_quantized(x: Value | Placeholder | Type) -> Value là lối tắt cho is_quantized_tensor_element_type(x).

  • is_type_name(x: Value | Placeholder | Type) -> Value. Dùng được cho mọi loại. Ví dụ: is_float(x) trả về true nếu xFloatType. Nếu x là một giá trị hoặc phần giữ chỗ, thì hàm này sẽ là lối tắt cho is_type_name(type(x)).

  • max_value(x: Type) -> Value trả về giá trị lớn nhất của TensorElementType. Nếu x không phải là TensorElementType, hãy trả về None.

  • min_value(x: Type) -> Value trả về giá trị tối thiểu có thể có là TensorElementType. Nếu x không phải là TensorElementType, hãy trả về None.

  • member_name(x: Value | Placeholder | Type) -> Any. Có sẵn cho mọi định nghĩa về thành phần member_name thuộc mọi loại. Ví dụ: tensor_element_type(x) trả về phần TensorElementType của TensorType tương ứng. Nếu x là một giá trị hoặc phần giữ chỗ, thì hàm này sẽ là lối tắt cho member_name(type(x)). Nếu x không phải là một loại có thành phần thích hợp hoặc một giá trị hay phần giữ chỗ thuộc loại như vậy, thì sẽ trả về None.

Xây dựng giá trị

  • operation_name(*xs: Value | Type) -> Value. Dùng được cho mọi thao tác. Ví dụ: add(lhs, rhs) nhận hai giá trị tensor lhsrhs, rồi trả về kết quả đánh giá hoạt động add với các dữ liệu đầu vào này. Đối với một số toán tử, chẳng hạn như broadcast_in_dim, các loại đầu ra của các toán tử này là "chịu tải", tức là cần thiết để đánh giá một hoạt động. Trong trường hợp này, hàm lấy các loại này làm đối số.

Hàm trên giá trị

  • Tất cả toán tử và hàm của Python đều có sẵn. Ví dụ: cả ký hiệu subscription (gói thuê bao) và sliceing (cắt lát) của Python đều có sẵn để lập chỉ mục thành tensor, tensor lượng tử hoá và bộ dữ liệu.

  • to_destination_type(x: Value, destination_type: Type) -> Value được xác định trên tensor và trả về giá trị chuyển đổi của x dựa trên type(x)destination_type như sau:

def to_destination_type(x: Value, destination_type: Type) -> Value:
  if type(x) == destination_type:
    return x

  if is_quantized(destination_type):
    if is_quantized(type(x)):
      return quantize(x, destination_type)
    assert is_float(type(x))
    return quantize(x, destination_type)

  if is_quantized(type(x)):
    assert destination_type = expressed_type(type(x))
    return dequantize(type(x))

  return convert(x, destination_type)

Có các cuộc thảo luận sớm về việc hợp nhất các thao tác convert, uniform_quantizeuniform_dequantize (#1576). Sau khi hợp nhất, chúng ta không cần hàm trên và có thể sử dụng tên thao tác cho convert.

  • is_nan(x: Value) -> Value được xác định trên tensor và trả về true nếu tất cả các phần tử của xNaN hoặc false. Nếu x không phải là một tensor, hãy trả về None.

  • is_sorted(x: Value) -> Value được xác định trên tensor và trả về true nếu các phần tử của x được sắp xếp theo thứ tự tăng dần so với thứ tự từ điển tăng dần của các chỉ mục hoặc false. Nếu x không phải là một tensor, hãy trả về None.

  • is_unique(x: Value) -> Value được xác định trên tensor và trả về true nếu x không có các phần tử trùng lặp hoặc false. Nếu x không phải là một tensor, hãy trả về None.

  • member_name(x: Value) -> Any được xác định cho mọi định nghĩa thành phần member_name của mọi giá trị. Ví dụ: real_part(x) trả về phần RealPart của ComplexConstant tương ứng. Nếu x không phải là một giá trị có thành phần thích hợp, thì trả về None.

  • same(x: Value) -> Value được xác định trên tensor và trả về true nếu các phần tử của x đều bằng nhau hoặc nếu không thì false. Nếu tensor không có các phần tử, thì tensor đó được tính là "tất cả bằng nhau", tức là hàm trả về true. Nếu x không phải là một tensor, hãy trả về None.

  • split(x: Value, num_results: Value, axis: Value) -> Value được xác định trên tensor và trả về các lát cắt num_results của x dọc theo trục axis. Nếu x không phải là một tensor hoặc dim(x, axis) % num_results != 0, trả về None.

Tính toán hình dạng

  • axes(x: Value | Placeholder | Type) -> Value là lối tắt cho range(rank(x)).

  • dim(x: Value | Placeholder | Type, axis: Value) -> Value là lối tắt cho shape(x)[axis].

  • dims(x: Value | Placeholder | Type, axes: List) -> List là lối tắt cho list(map(lambda axis: dim(x, axis), axes)).

  • index_space(x: Value | Placeholder | Type) -> Value được xác định trên tensor và trả về các chỉ mục size(x) cho TensorType tương ứng được sắp xếp theo thứ tự từ điển tăng dần, tức là [0, ..., 0], [0, ..., 1], ..., shape(x) - 1. Nếu x không phải là loại tensor, loại tensor lượng tử hoá, hoặc một giá trị hoặc phần giữ chỗ thuộc một trong những loại này, sẽ trả về None.

  • rank(x: Value | Placeholder | Type) -> Value là lối tắt cho size(shape(x)).

  • shape(x: Value | Placeholder | Type) -> Value được xác định trong phần "Hàm trên kiểu" thông qua member_name.

  • size(x: Value | Placeholder | Type) -> Value là lối tắt cho reduce(lambda x, y: x * y, shape(x)).

Các phép tính lượng tử hoá

  • def baseline_element_type(x: Value | Placeholder | Type) -> Type là phím tắt cho element_type(baseline_type(x)).

  • baseline_type được xác định trên các loại tensor và các loại tensor lượng tử hoá, đồng thời chuyển đổi chúng thành một "cơ sở", tức là một loại có cùng hình dạng nhưng có các tham số lượng tử hoá của loại phần tử được đặt lại về giá trị mặc định. Đây được dùng như một thủ thuật hữu ích để so sánh một cách thống nhất cả loại tensor và lượng tử hoá. Việc này thường xuyên cần được thực hiện. Đối với các loại lượng tử hoá, thuộc tính này cho phép so sánh các loại bỏ qua các tham số lượng tử hoá, nghĩa là shape, storage_type, expressed_type, storage_min, storage_maxquantization_dimension (đối với loại lượng tử hoá trên mỗi trục) phải khớp tất cả, nhưng scaleszero points có thể khác nhau.

def baseline_type(x: Value | Placeholder | Type) -> Type:
  if type(x) == TensorType:
    return x
  if type(x) == QuantizedTensorType:
    element_type = quantized_tensor_element_type(x)
    baseline_element_type = QuantizedTensorElementType(
      storage_type = storage_type(element_type),
      storage_min = storage_min(element_type),
      storage_max = storage_max(element_type),
      expressed_type = expressed_type(element_type),
      quantization_dimension = quantization_dimension(element_type),
      scales = [constant(1.0, expressed_type(element_type))] * dim(x, quantization_dimension(element_type)),
      zero_points = [constant(0, storage_type(element_type))] * dim(x, quantization_dimension(element_type)))
    return QuantizedTensorType(shape(x), baseline_element_type)
  if type(x) is not Type:
    return baseline_element_type(type(x))
  • dequantize được xác định trên các loại tensor lượng tử hoá và biến các tensor đó thành các loại tensor dấu phẩy động. Điều này xảy ra bằng cách chuyển đổi các phần tử lượng tử biểu thị các giá trị số nguyên của loại bộ nhớ thành các giá trị dấu phẩy động tương ứng của loại đã biểu thị bằng cách sử dụng điểm 0 và tỷ lệ liên kết với loại phần tử lượng tử hoá.
def compute_zero_points(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(zero_point(quantized_type), storage_type(quantized_type)), [], result_type)
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      zero_points[i] = zero_points(quantized_type)[i[d]]
    return zero_points

def compute_scales(quantized_type, result_type):
  if is_per_tensor_quantized(quantized_type):
    return broadcast_in_dim(constant(scale(quantized_type), expressed_type(quantized_type)), [],
            type(result_type))
  if is_per_axis_quantized(quantized_type):
    for i in index_space(result_type):
      d = quantization_dimension(quantized_type)
      scales[i] = scales(quantized_type)[i[d]]
    return scales

def dequantize(x: Value) -> Value:
  assert is_quantized(x)
  x_storage = bitcast_convert(x, storage_type(x))
  x_storage_sub = x_storage - compute_zero_points(type(x), type(x_storage))
  x_expressed_sub = convert(x_storage_sub, expressed_type(x))
  return x_expressed_sub * compute_scales(type(x), type(x_expressed_sub))
  • quantize được xác định trên các loại tensor dấu phẩy động và biến chúng thành các loại tensor lượng tử hoá. Điều này xảy ra bằng cách chuyển đổi các giá trị dấu phẩy động của loại được biểu thị thành các giá trị số nguyên tương ứng của loại bộ nhớ bằng cách sử dụng điểm 0 và tỷ lệ liên kết với loại phần tử lượng tử hoá.
def quantize(x: Value, type: Type) -> Value:
  assert is_float(x) and is_quantized(type)
  x_expressed_rounded = round_nearest_even(x / compute_scales(type, type(x)))
  x_storage_rounded = convert(x_expressed_rounded, storage_type(type))
  x_storage_add = x_storage_rounded + compute_zero_points(type, type(x_storage_rounded))
  x_storage = clamp(storage_min(type), x_storage_add, storage_max(type))
  return bitcast_convert(x_storage, type)
  • dequantize_op_quantize dùng để chỉ định các phép tính toán phần tử trên tensor lượng tử hoá. Trình phân tích cú pháp này sẽ khử lượng tử, tức là chuyển các phần tử lượng tử thành các loại được biểu thị, sau đó thực hiện một thao tác, sau đó lượng tử hoá, tức là chuyển kết quả trở về loại bộ nhớ của chúng. Hiện tại, hàm này chỉ hoạt động đối với quá trình lượng tử hoá mỗi tensor. Quá trình lượng tử hoá trên mỗi trục đang diễn ra (#1574).
def dequantize_op_quantize(op, *inputs_and_output_type):
  inputs = inputs_and_output_type[:-1]
  output_type = inputs_and_output_type[-1]

  float_inputs = map(dequantize, inputs)
  float_result = op(*float_inputs)
  return quantize(float_result, output_type)

def dequantize_batch_norm_grad_or_training_quantize(op, *inputs_and_output_types):
  inputs = inputs_and_output_type[:-3]
  float_inputs = map(dequantize, inputs)
  float_results = op(*float_inputs)
  return map(quantize, float_results, inputs_and_output_type[-3:])

def dequantize_compare(lhs, rhs, comparison_direction):
  float_lhs = dequantize(lhs)
  float_rhs = dequantize(rhs)
  return compare(float_lhs, float_rhs, comparison_direction, FLOAT)

def dequantize_select_quantize(pred, on_true, on_false, output_type):
  float_on_true = dequantize(on_true)
  float_on_false = dequantize(on_false)
  float_result = select(pred, float_on_true, float_on_false)
  return quantize(float_result, output_type)

Tính toán lưới

  • cross_partition(replica_groups: Value) -> Value. Xem phần "cross_replica" ở trên.

  • cross_replica(replica_groups: Value) -> Value. Xem phần "cross_replica" ở trên.

  • cross_replica_and_partition(replica_groups: Value) -> Value. Xem phần "cross_replica_and_partition" ở trên.

  • flattened_ids(replica_groups: Value) -> Value. Xem phần "flattened_ids" ở trên.