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

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

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

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

Để xem quy cách của một bản phát hành trước của StableHLO, hãy mở kho lưu trữ tại bản phát hành được gắn thẻ mà bạn quan tâm. Ví dụ: Quy cách StableHLO phiên bản 0.19.0. Để xem những thay đổi xảy ra ở mỗi lần tăng phiên bản phụ của StableHLO, hãy tham khảo nhật ký phiên bản trong VhloDialect.td.

Chương trình

Program ::= {Func}

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

func.func @main(
  %image: tensor<28x28xf32>,
  %weights: tensor<784x10xf32>,
  %bias: tensor<1x10xf32>
) -> tensor<1x10xf32> {
  %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
  %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
  %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
  %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
  "func.return"(%4): (tensor<1x10xf32>) -> ()
}

Hàm

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

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

Giá trị nhận dạng

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

Giá trị nhận dạng StableHLO tương tự như giá trị nhận dạng trong nhiều ngôn ngữ lập trình, với 2 đặc điểm riêng biệt: 1) tất cả giá trị nhận dạng đều có ký 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ể hoàn toàn là số để đơn giản hoá việc tạo chương trình StableHLO.

Loại

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

Các loại StableHLO được phân loại thành các loại giá trị (còn được gọi là các loại hạng nhất) đại diện cho các giá trị StableHLO và các loại không phải giá trị mô tả các phần tử 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, điểm đặc biệt chính là bản chất dành riêng cho miền của StableHLO, dẫn đến một số kết quả bất thường (ví dụ: các loại vô hướng không phải là các loại giá trị).

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

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

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

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

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

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

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

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

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

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

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

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

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

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

TokenType ::= 'token'

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

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

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

Bạn có thể phân bổ các vùng đệm bằng cách sử dụng custom_call thành CreateBuffer hoặc Pin và huỷ phân bổ thông qua custom_call thành Unpin. Chỉ các thao tác custom_call mới có thể đọc và ghi nội dung bên trong các vùng đệm. Hãy xem custom_call để biết thêm thông tin chi tiết.

Các loại bộ biểu thị các bộ, tức là các danh sách không đồng nhất. Bộ giá trị là một tính năng cũ chỉ tồn tại để tương thích với HLO. Trong HLO, các bộ được dùng để biểu thị đầu vào và đầu ra có thể thay đổi. Trong StableHLO, các đầu vào và đầu ra có số lượng thay đổi được hỗ trợ một cách tự nhiên và việc sử dụng bộ dữ liệu duy nhất trong StableHLO là để thể hiện toàn diện ABI HLO, trong đó chẳng hạn như T, tuple<T> và 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 thay đổi ABI HLO. Điều này có thể cho phép chúng tôi xoá các loại bộ dữ liệu khỏi StableHLO (#598).

TensorElementType ::= BooleanType | IntegerType | FloatType | ComplexType
BooleanType ::= 'i1'
IntegerType ::= SignedIntegerType | UnsignedIntegerType
SignedIntegerType ::= 'si2' | 'si4' | 'si8' | 'si16' | 'si32' | 'si64'
UnsignedIntegerType ::= 'ui2' | 'ui4' | 'ui8' | 'ui16' | 'ui32' | 'ui64'
FloatType ::= 'f4E2M1FN' | 'f6E2M3FN' | 'f6E3M2FN' | 'f8E3M4' | 'f8E4M3'
            | 'f8E4M3FN' | 'f8E4M3FNUZ' | 'f8E4M3B11FNUZ' | 'f8E5M2'
            | 'f8E5M2FNUZ' | 'f8E8M0FNU' | 'bf16' | 'f16' | 'f32' | 'f64'
TensorFloat32 ::= 'tf32'
ComplexType ::= 'complex' '<' ComplexElementType '>'
ComplexElementType ::= 'f32' | 'f64'

Các loại phần tử đại diện cho các phần tử của các loại tensor. Không giống như trong nhiều ngôn ngữ lập trình, các loại này không phải là loại hạng nhất 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 những loại này (do đó, theo thông lệ, bạn nên biểu thị các giá trị vô hướng thuộc loại T bằng các giá trị tensor 0 chiều thuộc loại tensor<T>).

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

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

Các loại hàm biểu thị cả hàm được đặt tên và hàm ẩn danh. Chúng có các loại đầu vào (danh sách các loại ở phía bên trái của ->) và các loại đầu ra (danh sách các loại ở phía bên phải của ->). Trong nhiều ngôn ngữ lập trình, các loại hàm là loại đầ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, loại chuỗi không phải là loại dữ liệu hạng nhất trong StableHLO và chỉ được dùng để chỉ định siêu dữ liệu tĩnh cho các phần tử chương trình.

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

Các thao tác StableHLO (còn được gọi là ops) đại diện cho một tập hợp khép kín gồm các thao tác cấp cao trong các mô hình học máy. Như đã thảo luận ở trên, cú pháp StableHLO chịu ảnh hưởng lớn của MLIR. Đây không nhất thiết là lựa chọn thay thế tiện dụng nhất, nhưng có thể nói là phù hợp nhất với mục tiêu của StableHLO là tạo ra khả năng tương tác cao hơn giữa các khung ML 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 được gọi là ops) 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 từ viết tắt giúp xác định duy nhất một trong các thao tác được hỗ trợ. Xem bên dưới để biết danh sách đầy đủ tất cả các thao tác được hỗ trợ.

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

Các thao tác sử dụng đầu vào và tạo ra đầu ra. Đầu vào được phân loại thành các giá trị đầu vào (được tính toán trong quá trình thực thi), hàm đầu vào (được cung cấp tĩnh, vì trong StableHLO, các hàm không phải là giá trị lớp thứ nhất) và các thuộc tính đầu vào (cũng được cung cấp tĩnh). Loại đầu vào và đầu ra mà một thao tác sử dụng và tạo ra phụ thuộc vào từ viết tắt của thao tác đó. 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, thao tác select_and_scatter sẽ sử dụng 3 giá trị đầu vào, 2 hàm đầu vào và 3 thuộc tính đầu vào.

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

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

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

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

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

Hiện tại, các chương trình StableHLO trong 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 những thuộc tính này vào opset 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, bạn có thể xem danh sách các thuộc tính này tại đây:

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

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

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

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

%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 chữ và một loại, cùng nhau biểu thị một giá trị StableHLO. Nhìn chung, kiểu này là một phần của cú pháp hằng số, ngoại trừ khi kiểu này không rõ ràng (ví dụ: hằng số boolean có kiểu i1 không rõ ràng, trong khi hằng số nguyên có thể có nhiều kiểu có thể có).

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

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

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

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

• (C1) is_wellformed(integer_literal, integer_type).

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

Hằng số dấu phẩy động biểu thị các giá trị dấu phẩy động thông qua các chuỗi sử dụng ký hiệu thập phân hoặc 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 ràng buộc sau:

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

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

Hằng số phức biểu thị các giá trị phức bằng cách sử dụng danh sách gồm phần thực (xuất hiện đầu tiên) và phần ảo (xuất hiện thứ hai). Ví dụ: (1.0, 0.0) : complex<f32> biểu thị 1.0 + 0.0i và (0.0, 1.0) : complex<f32> biểu thị 0.0 + 1.0i. Thứ tự mà các phần này được lưu trữ trong bộ nhớ là do quá trình triển khai xác định. Hằng số phức có các ràng buộc sau:

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

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

Hằng số tensor biểu thị các giá trị tensor bằng cách sử dụng các danh sách lồng nhau được chỉ định thông qua ký pháp NumPy. Ví dụ: dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> biểu thị một giá trị tensor có ánh xạ sau đây từ chỉ mục đến các phần tử: {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5, {1, 2} => 6. Thứ tự mà các phần tử này được lưu trữ trong bộ nhớ sau đó sẽ do quá trình triển khai xác định. Hằng số Tensor có cá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, hãy nói false.

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

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

• (C1) has_syntax(quantized_tensor_literal, storage_type(quantized_tensor_type)).
• (C2) has_shape(quantized_tensor_literal, shape(quantized_tensor_type)).

StringConstant  ::= StringLiteral
StringLiteral   ::= '"' {stringCharacter | escapeSequence} '"'
stringCharacter ::= all ASCII characters except '\00', '\01', ... '\1f' and '"'
escapeSequence  ::= '\' ('"' | '\' | 'n' | 't' | (hexadecimalDigit hexadecimalDigit))

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

Ops

Giá trị tuyệt đối

Ngữ nghĩa

Thực hiện phép toán abs theo phần tử trên tensor operand và tạo ra một tensor result. Tuỳ thuộc vào loại phần tử, hãy 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 được lượng tử hoá: dequantize_op_quantize(abs, operand, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

thêm

Ngữ nghĩa

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

• Đối với giá trị boolean: OR logic.
• Đối với số nguyên: phép 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 số phức.
• Đối với các loại được lượng tử hoá: dequantize_op_quantize(add, lhs, rhs, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• Nếu thao tác sử dụng các tensor không được định lượng:
  • (C1) type(lhs) = type(rhs) = type(result).
• Nếu thao tác sử dụng các tensor được định lượng:
  • (C2) is_quantized(lhs) and is_quantized(rhs) and is_quantized(result).
  • (C3) storage_type(lhs) = storage_type(rhs) = storage_type(result).
  • (C4) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C5) (is_per_axis_quantized(lhs) or is_per_axis_quantized(rhs)) = is_per_axis_quantized(result).
  • (C6) Nếu is_per_axis_quantized(lhs), thì quantization_dimension(lhs) = quantization_dimension(result).
  • (C7) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) = quantization_dimension(result).

Ví dụ

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

 Các ví dụ khác

after_all

Ngữ nghĩa

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

Thông tin đầu vào

Kết quả đầu ra

Ví dụ

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

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

Thao tác này 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.
• results...@process = concatenate(operands...@process, all_gather_dim) cho tất cả process trong process_group.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [[1, 2], [3, 4]]
// %operand0@(1, 0): [[5, 6], [7, 8]]
// %operand1@(0, 0): [[11, 12], [13, 14]]
// %operand1@(1, 0): [[15, 16], [17, 18]]
%result:2 = "stablehlo.all_gather"(%operand0, %operand1) {
  all_gather_dim = 1 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<2x2xi64>, tensor<2x2xi64>) -> (tensor<2x4xi64>, tensor<2x4xi64>)
// %result0@(0, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result0@(1, 0): [[1, 2, 5, 6], [3, 4, 7, 8]]
// %result1@(0, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]
// %result1@(1, 0): [[11, 12, 15, 16], [13, 14, 17, 18]]

 Các 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 operands từ mỗi quy trình và tạo ra các tensor results.

Thao tác này 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:

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand0@(0, 0): [1, 2, 3, 4]
// %operand0@(1, 0): [5, 6, 7, 8]
// %operand1@(0, 0): [9, 10, 11, 12]
// %operand1@(1, 0): [13, 14, 15, 16]
%result:2 = "stablehlo.all_reduce"(%operand0, %operand0) ({
  ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
    %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
    "stablehlo.return"(%0) : (tensor<i64>) -> ()
}) {
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
  // channel_id = 0
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
  // use_global_device_ids = false
} : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
// %result0@(0, 0): [6, 8, 10, 12]
// %result0@(1, 0): [6, 8, 10, 12]
// %result1@(0, 0): [22, 24, 26, 28]
// %result1@(1, 0): [22, 24, 26, 28]

 Các ví dụ khác

all_to_all

Ngữ nghĩa

Trong mỗi nhóm quy trình trong lưới quy trình StableHLO, hãy chia các giá trị của tensor operands dọc theo split_dimension thành các phần, phân tán các phần đã chia giữa các quy trình, nối các phần đã phân tán dọc theo concat_dimension và tạo ra các tensor results. Thao tác này 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(operands...@sender, split_count, split_dimension) cho tất cả sender trong process_group.
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] nơi receiver_index = process_group.index(receiver).
• results...@process = concatenate(scattered_parts...@process, concat_dimension).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// num_replicas: 2
// num_partitions: 1
// %operand1@(0, 0): [[1, 2, 3, 4],
//                    [5, 6, 7, 8]]
// %operand1@(1, 0): [[9, 10, 11, 12],
//                    [13, 14, 15, 16]]
// %operand2@(0, 0): [[17, 18, 19, 20],
//                    [21, 22, 23, 24]]
// %operand2@(1, 0): [[25, 26, 27, 28],
//                    [29, 30, 31, 32]]
%result:2 = "stablehlo.all_to_all"(%operand1, %operand2) {
  split_dimension = 1 : i64,
  concat_dimension = 0 : i64,
  split_count = 2 : i64,
  replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
  // channel_id = 0
} : (tensor<2x4xi64>, tensor<2x4xi64>) -> (tensor<4x2xi64>, tensor<4x2xi64>)
// %result#0@(0, 0): [[1, 2], [5, 6], [9, 10], [13, 14]]
// %result#0@(1, 0): [[3, 4], [7, 8], [11, 12], [15, 16]]
// %result#1@(0, 0): [[17, 18], [21, 22], [25, 26], [29, 30]]
// %result#1@(1, 0): [[19, 20], [23, 24], [27, 28], [31, 32]]

 Các ví dụ khác

và

Ngữ nghĩa

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

• Đối với giá trị boolean: phép toán AND logic.
• Đối với số nguyên: bitwise AND.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

atan2

Ngữ nghĩa

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %lhs: [0.0, 1.0, -1.0]
// %rhs: [0.0, 0.0, 0.0]
%result = "stablehlo.atan2"(%lhs, %rhs) : (tensor<3xf64>, tensor<3xf64>) -> tensor<3xf64>
// %result: [0.0, 1.57079637, -1.57079637] // [0.0, pi/2, -pi/2]

 Các ví dụ khác

batch_norm_grad

Ngữ nghĩa

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

def compute_sum(operand, feature_index):
  (sum,) = reduce(
      inputs=[operand],
      init_values=[constant(0, element_type(operand))],
      dimensions=[i for i in range(rank(operand)) if i != feature_index],
      body=lambda x, y: add(x, y))
  return sum

def compute_mean(operand, feature_index):
  sum = compute_sum(operand, feature_index)
  divisor = constant(size(operand) / dim(operand, feature_index),
                     element_type(operand))
  divisor_bcast = broadcast_in_dim(divisor, [], type(sum))
  return divide(sum, divisor_bcast)

def batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index):
  # Broadcast inputs to type(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance`
  # Intermediate values will be useful for computing gradients
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)

  # Use the implementation from batchnorm_expander.cc in XLA
  # Temporary variables have exactly the same names as in the C++ code
  elements_per_feature = broadcast_in_dim(
      constant(divide(size(operand), dim(operand, feature_index)),
               element_type(grad_output)),
      [], type(operand))
  i1 = multiply(grad_output, elements_per_feature)
  i2 = broadcast_in_dim(
      compute_sum(grad_output, feature_index), [feature_index], type(operand))
  i3 = broadcast_in_dim(
      compute_sum(multiply(grad_output, centered_operand), feature_index),
      [feature_index], type(operand))
  i4 = multiply(i3, centered_operand)
  i5 = divide(i4, add(variance_bcast, epsilon_bcast))
  i6 = subtract(subtract(i1, i2), i5)

  grad_operand =
      multiply(divide(divide(scale_bcast, stddev), elements_per_feature), i6)
  grad_scale =
      compute_sum(multiply(grad_output, normalized_operand), feature_index)
  grad_offset = compute_sum(grad_output, feature_index)

  return grad_operand, grad_scale, grad_offset

Đối với các loại được định lượng, hãy thực hiện dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, mean, variance, grad_output: batch_norm_grad(operand, scale, mean, variance, grad_output, epsilon, feature_index), operand, scale, mean, variance, grad_output, type(grad_operand), type(grad_scale), type(feature_index)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

batch_norm_inference

Ngữ nghĩa

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

def batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index):
  # Broadcast inputs to shape(operand)
  scale_bcast = broadcast_in_dim(scale, [feature_index], type(operand))
  offset_bcast = broadcast_in_dim(offset, [feature_index], type(operand))
  mean_bcast = broadcast_in_dim(mean, [feature_index], type(operand))
  variance_bcast = broadcast_in_dim(variance, [feature_index], type(operand))
  epsilon_bcast = broadcast_in_dim(constant(epsilon, element_type(operand)), [],
                                   type(operand))

  # Perform normalization using the provided `mean` and `variance` instead of
  # computing them like `batch_norm_training` does.
  centered_operand = subtract(operand, mean_bcast)
  stddev = sqrt(add(variance_bcast, epsilon_bcast))
  normalized_operand = divide(centered_operand, stddev)
  return add(multiply(scale_bcast, normalized_operand), offset_bcast)

Đối với các loại được định lượng, hãy thực hiện dequantize_op_quantize(lambda operand, scale, offset, mean, variance: batch_norm_inference(operand, scale, offset, mean, variance, epsilon, feature_index), operand, scale, offset, mean, variance, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, mean, variance và result có cùng baseline_element_type.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(mean) = dim(operand, feature_index).
• (C6) size(variance) = dim(operand, feature_index).
• (C7) baseline_type(operand) = baseline_type(result).

Ví dụ

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
// %mean: [2.0, 3.0]
// %variance: [1.0, 1.0]
%result = "stablehlo.batch_norm_inference"(%operand, %scale, %offset, %mean, %variance) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>, tensor<2xf64>) -> tensor<2x2x2xf64>
// %result: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]

batch_norm_training

Ngữ nghĩa

Tính toán giá trị trung bình và phương sai trên tất cả các phương diện, ngoại trừ phương diện feature_index và chuẩn hoá tenxơ operand tạo ra các tenxơ output, batch_mean và batch_var. Cụ thể hơn, bạn có thể biểu thị thao tác này dưới dạng một thao tác phân tách thành các thao tác 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 được định lượng, hãy thực hiện dequantize_batch_norm_grad_or_training_quantize(lambda operand, scale, offset: batch_norm_training(operand, scale, offset, epsilon, feature_index), operand, scale, offset, type(output), type(batch_mean), type(batch_var)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) 0 <= feature_index < rank(operand).
• (C2) operand, scale, offset, batch_mean, batch_var và output có cùng baseline_element_type.
• (C3) size(scale) = dim(operand, feature_index).
• (C4) size(offset) = dim(operand, feature_index).
• (C5) size(batch_mean) = dim(operand, feature_index).
• (C6) size(batch_var) = dim(operand, feature_index).
• (C7) baseline_type(output) = baseline_type(operand).

Ví dụ

// %operand: [
//            [[1.0, 2.0], [3.0, 4.0]],
//            [[3.0, 4.0], [1.0, 2.0]]
//           ]
// %scale: [1.0, 1.0]
// %offset: [1.0, 1.0]
%output, %batch_mean, %batch_var = "stablehlo.batch_norm_training"(%operand, %scale, %offset) {
  epsilon = 0.0 : f32,
  feature_index = 2 : i64
} : (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>) ->
    (tensor<2x2x2xf64>, tensor<2xf64>, tensor<2xf64>)
// %output: [
//           [[0.0, 0.0], [2.0, 2.0]],
//           [[2.0, 2.0], [0.0, 0.0]]
//          ]
// %batch_mean: [2.0, 3.0]
// %batch_var: [1.0, 1.0]

bitcast_convert

Ngữ nghĩa

Thực hiện một thao tác bitcast trên tenxơ operand và tạo ra một tenxơ result, trong đó các bit của toàn bộ tenxơ operand được diễn giải lại bằng cách sử dụng kiểu của tenxơ result.

Cụ thể hơn, cho E = element_type(operand), E' = element_type(result) và 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ề biểu thị trong bộ nhớ của một giá trị nhất định và hành vi của giá trị này được xác định theo cách triển khai vì biểu thị chính xác của các tensor được xác định theo cách triển khai và biểu thị chính xác của các loại phần tử cũng được xác định theo cách triển khai.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) Cho E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) và R = rank(operand):
  • Nếu num_bits(E') = num_bits(E), shape(result) = shape(operand).
  • Nếu 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 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

 Các ví dụ khác

broadcast_in_dim

Ngữ nghĩa

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %operand: [
//            [1, 2, 3]
//           ]
%result = "stablehlo.broadcast_in_dim"(%operand) {
  broadcast_dimensions = array<i64: 2, 1>
} : (tensor<1x3xi32>) -> tensor<2x3x2xi32>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Các ví dụ khác

ốp lưng

Ngữ nghĩa

Tạo ra kết quả đầu ra 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. Cụ thể hơn, result = selected_branch() trong đó:

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %index: -1
// %result_branch0: [0, 0]
// %result_branch1: [1, 1]
%result0, %result1 = "stablehlo.case"(%index) ({
  "stablehlo.return"(%result_branch0, %result_branch0) : (tensor<2xi64>, tensor<2xi64>) -> ()
}, {
  "stablehlo.return"(%result_branch1, %result_branch1) : (tensor<2xi64>, tensor<2xi64>) -> ()
}) : (tensor<i32>) -> (tensor<2xi64>, tensor<2xi64>)
// %result0: [1, 1]
// %result1: [1, 1]

 Các ví dụ khác

cbrt

Ngữ nghĩa

Thực hiện phép toán căn bậc ba 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: rootn(x, 3) từ IEEE-754.
• Đối với số phức: căn bậc ba của số phức.
• Đối với các loại được định lượng: dequantize_op_quantize(cbrt, operand, type(result))

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %operand: [0.0, 1.0, 8.0, 27.0]
%result = "stablehlo.cbrt"(%operand) : (tensor<4xf64>) -> tensor<4xf64>
// %result: [0.0, 1.0, 2.0, 3.0]

 Các ví dụ khác

ceil

Ngữ nghĩa

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %operand: [-0.8166, -0.2530, 0.2530, 0.8166, 2.0]
%result = "stablehlo.ceil"(%operand) : (tensor<5xf32>) -> tensor<5xf32>
// %result: [-0.0, -0.0, 1.0, 1.0, 2.0]

 Các ví dụ khác

cholesky

Ngữ nghĩa

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

Cụ thể hơn, đối với tất cả i trong index_space(result), result[i0, ..., iR-3, :, :] là một phép phân tích Cholesky của a[i0, ..., iR-3, :, :], dưới dạng ma trận tam giác dưới (nếu lower là true) hoặc ma trận tam giác trên (nếu lower là false). 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 cách triển khai.

Nếu có 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 được định lượng, hãy thực hiện dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

kẹp

Ngữ nghĩa

Kẹp mọi phần tử của tensor operand trong khoảng từ giá trị tối thiểu đến tối đa và tạo ra tensor result. Cụ thể hơn, result[result_index] = minimum(maximum(operand[result_index], min_element), max_element), trong đó min_element = rank(min) = 0 ? min[] : min[result_index], max_element = rank(max) = 0 ? max[] : max[result_index]. Đối với các loại được lượng tử hoá, hãy thực hiện dequantize_op_quantize(clamp, min, operand, max, type(result)).

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

collective_broadcast

Ngữ nghĩa

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

Thao tác này 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 tính theo công thức:

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// num_replicas: 4
// num_partitions: 1
// %operand@(0, 0): [[1, 2]]
// %operand@(1, 0): [[3, 4]]
// %operand@(2, 0): [[5, 6]]
// %operand@(3, 0): [[7, 8]]
%result = "stablehlo.collective_broadcast"(%operand) {
  replica_groups = dense<[[2, 1]]> : tensor<1x2xi64>,
  channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
} : (tensor1x2xi64>) -> tensor<1x2xi64>
// %result@(0, 0): [[0, 0]]
// %result@(1, 0): [[5, 6]]
// %result@(2, 0): [[5, 6]]
// %result@(3, 0): [[0, 0]]

collective_permute

Ngữ nghĩa

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

Thao tác này 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 tính theo công thức:

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) dim(source_target_pairs, 1) = 2.
• (C2) is_unique(source_target_pairs[:, 0]).
• (C3) is_unique(source_target_pairs[:, 1]).
• (C4) 0 <= source_target_pairs < N, trong đó N được xác định là:
  • num_replicas nếu bạn sử dụng cross_replica.
  • num_partitions nếu bạn 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]]

 Các ví dụ khác

so sánh

Ngữ nghĩa

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

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

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

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

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

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

Đối với các kiểu phần tử dấu phẩy động có compare_type = TOTALORDER, thao tác này sử dụng tổ hợp các thao tác totalOrder và compareQuietEqual từ IEEE-754.

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

Đối với các loại được định lượng, hãy thực hiện dequantize_compare(lhs, rhs, comparison_direction).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %lhs: [1.0, 3.0]
// %rhs: [1.1, 2.9]
%result = "stablehlo.compare"(%lhs, %rhs) {
  comparison_direction = #stablehlo<comparison_direction LT>,
  compare_type = #stablehlo<comparison_type FLOAT>
} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>
// %result: [true, false]

 Các ví dụ khác

phức tạp

Ngữ nghĩa

Thực hiện quá trình 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à giá trị ảo, lhs và rhs, đồng thời tạo ra một tenxơ result.

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %lhs: [1.0, 3.0]
// %rhs: [2.0, 4.0]
%result = "stablehlo.complex"(%lhs, %rhs) : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xcomplex<f64>>
// %result: [(1.0, 2.0), (3.0, 4.0)]

 Các ví dụ khác

tổng hợp

Ngữ nghĩa

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) is_namespaced_op_name(name)
• (C2) is_defined_in_parent_scope(decomposition)
• (C3) types(inputs...) == input_types(decomposition)
• (C4) types(results...) == output_types(decomposition)

Ví dụ

%results = "stablehlo.composite"(%input0, %input1) {
  name = "my_namespace.my_op",
  composite_attributes = {
    my_attribute = "my_value"
  },
  decomposition = @my_op,
  version = 1 : i32
} : (tensor<f32>, tensor<f32>) -> tensor<f32>

 Các ví dụ khác

nối

Ngữ nghĩa

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) same(element_type(inputs...)).
• (C2) same(shape(inputs...)) ngoại trừ dim(inputs..., dimension).
• (C3) 0 < size(inputs).
• (C4) 0 <= dimension < rank(inputs[0]).
• (C5) element_type(result) = element_type(inputs[0]).
• (C6) shape(result) = shape(inputs[0]) ngoại trừ:
  • dim(result, dimension) = dim(inputs[0], dimension) + ....

Ví dụ

// %input0: [[1, 2], [3, 4], [5, 6]]
// %input1: [[7, 8]]
%result = "stablehlo.concatenate"(%input0, %input1) {
  dimension = 0 : i64
} : (tensor<3x2xi64>, tensor<1x2xi64>) -> tensor<4x2xi64>
// %result: [[1, 2], [3, 4], [5, 6], [7, 8]]

 Các ví dụ khác

hằng số

Ngữ nghĩa

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

%output = "stablehlo.constant"() {
  value = dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>
} : () -> tensor<2x2xf32>
// %output: [[0.0, 1.0], [2.0, 3.0]]

 Các ví dụ khác

chuyển đổi

Ngữ nghĩa

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

Đối với các lượt chuyển đổi boolean-to-any-supported-type, giá trị false được chuyển đổi thành 0 và giá trị true được chuyển đổi thành 1. Đối với các lượt chuyển đổi any-supported-type-to-boolean, giá trị 0 sẽ được chuyển đổi thành false và các giá trị khác 0 sẽ được chuyển đổi thành true. Hãy xem bên dưới để biết cách hoạt động này đối với các loại phức tạp.

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

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

Quá trình chuyển đổi liên quan đến số phức sang số phức tuân theo cùng một hành vi của quá trình chuyển đổi số thực sang số thực để chuyển đổi các phần thực và phần ảo.

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %operand: [-1, 0, 1]
%result = "stablehlo.convert"(%operand) : (tensor<3xi64>) -> tensor<3xcomplex<f64>>
// %result: [(-1.0, 0.0), (0.0, 0.0), (1.0, 0.0)]

 Các ví dụ khác

tích chập

Ngữ nghĩa

Tính tích vô hướng giữa các cửa sổ của lhs và các lá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ừ lhs và rhs bằng một ví dụ cụ thể.

Cụ thể hơn, hãy xem xét việc điều chỉnh lại các đầu vào sau đây theo lhs để có thể biểu thị các cửa sổ 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 điều chỉnh khung hình này sử dụng các hàm trợ giúp sau:

• lhs_shape(n, hw, c) = permute([n] + hw + [c], [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]).
• result_shape(n1, hw, c1) = permute([n1] + hw + [c1], [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]).
• permute([j0, j1, ..., jR-1], permutation) = [i0, i1, ..., iR-1] nơi j[d] = i[permutation[d]].

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

• padding_value = constant(0, element_type(lhs)).
• padded_lhs = pad(lhs, padding_value, lhs_padding[:, 0], lhs_padding[:, 1], lhs_base_dilations - 1).
• lhs_window_start = lhs_shape(0, output_spatial_index, 0) * lhs_window_strides.
• lhs_window = slice(padded_lhs, lhs_window_start, lhs_window_start + lhs_window_dimensions, lhs_window_dilations).
• reversed_lhs_window = reverse(lhs_window, [input_spatial_dimensions[dim] for dim in range(size(window_reversal)) if window_reversal[dim] = true]). Tính năng này có vẻ không được dùng, nên trong tương lai, chúng tôi dự định sẽ xoá tính năng này (#1181).
• dot_product = dot_general(reversed_lhs_window, rhs, lhs_batching_dimensions=[], lhs_contracting_dimensions=input_spatial_dimensions + [input_feature_dimension], rhs_batching_dimensions=[], rhs_contracting_dimensions=kernel_spatial_dimensions + [kernel_input_feature_dimension]).

Nếu feature_group_count > 1:

• lhses = split(lhs, feature_group_count, input_feature_dimension).
• rhses = split(rhs, feature_group_count, kernel_output_feature_dimension).
• results... = convolution(lhses..., rhses..., ..., feature_group_count=1, ...).
• result = concatenate(results, output_feature_dimension).

Nếu batch_group_count > 1:

• lhses = split(lhs, batch_group_count, input_batch_dimension).
• rhses = split(rhs, batch_group_count, kernel_output_feature_dimension).
• results... = convolution(lhses..., rhses..., ..., batch_group_count=1, ...).
• result = concatenate(results, output_feature_dimension).

Đối với các loại được lượng tử hoá, hãy thực hiện dequantize_op_quantize( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs, type(result)).

Đối với các loại được lượng tử hoá kết hợp, hãy thực hiện hybrid_dequantize_then_op( lambda lhs, rhs: convolution(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, window_reversal, input_batch_dimension, input_feature_dimension, input_spatial_dimensions, kernel_input_feature_dimension, kernel_output_feature_dimension, kernel_spatial_dimensions, output_batch_dimension, output_feature_dimension, output_spatial_dimensions, feature_group_count, batch_group_count, precision_config), lhs, rhs).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) N = rank(lhs) = rank(rhs).
• (C2) size(window_strides) = N - 2.
• (C3) 0 < window_strides.
• (C4) shape(padding) = [N - 2, 2].
• (C5) size(lhs_dilation) = N - 2.
• (C6) 0 < lhs_dilation.
• (C7) size(rhs_dilation) = N - 2.
• (C8) 0 < rhs_dilation.
• (C9) size(window_reversal) = N - 2.
• (C10) dim(lhs, input_batch_dimension) % batch_group_count = 0.
• (C11) dim(lhs, input_feature_dimension) % feature_group_count = 0.
• (C12) size(input_spatial_dimensions) = N - 2.
• (C13) Cho input_dimensions = [input_batch_dimension] + input_spatial_dimensions + [input_feature_dimension]:
  • is_unique(input_dimensions).
  • 0 <= input_dimensions < N.
• (C14) dim(rhs, kernel_input_feature_dimension) = dim(lhs, input_feature_dimension) / feature_group_count.
• (C15) dim(rhs, kernel_output_feature_dimension) % batch_group_count = 0.
• (C16) dim(rhs, kernel_output_feature_dimension) % feature_group_count = 0.
• (C17) size(kernel_spatial_dimensions) = N - 2.
• (C18) Điều kiện kernel_dimensions = kernel_spatial_dimensions + [kernel_input_feature_dimension] + [kernel_output_feature_dimension]:
  • is_unique(kernel_dimensions).
  • 0 <= kernel_dimensions < N.
• (C19) size(output_spatial_dimensions) = N - 2.
• (C20) Cho output_dimensions = [output_batch_dimension] + output_spatial_dimensions + [output_feature_dimension]:
  • is_unique(output_dimensions).
  • 0 <= output_dimensions < N.
• (C21) 0 < feature_group_count.
• (C22) 0 < batch_group_count.
• (C23) feature_group_count = 1 or batch_group_count = 1.
• (C24) size(precision_config) = 2.
• (C25) dim(result, result_dim) được xác định là:
  • dim(lhs, input_batch_dimension) / batch_group_count nếu result_dim = output_batch_dimension.
  • dim(rhs, kernel_output_feature_dimension) nếu result_dim = output_feature_dimension.
  • num_windows nếu không, trong đó:
  • output_spatial_dimensions[spatial_dim] = result_dim.
  • lhs_dim = input_spatial_dimensions[spatial_dim].
  • rhs_dim = kernel_spatial_dimensions[spatial_dim].
  • dilated_input_shape[lhs_dim] = dim(lhs, lhs_dim) = 0 ? 0 : (dim(lhs, lhs_dim) - 1) * lhs_dilation[spatial_dim] + 1.
  • padded_input_shape[lhs_dim] = padding[spatial_dim, 0] + dilated_input_shape[lhs_dim] + padding[spatial_dim, 1].
  • dilated_window_shape[lhs_dim] = dim(rhs, rhs_dim) = 0 ? 0 : (dim(rhs, rhs_dim) - 1) * rhs_dilation[spatial_dim] + 1.
  • is_empty_window[lhs_dim] = padded_input_shape[lhs_dim] = 0 || dilated_window_shape[lhs_dim] > padded_input_shape[lhs_dim].
  • num_windows = is_empty_window[lhs_dim] ? 0 : floor((padded_input_shape[lhs_dim] - dilated_window_shape[lhs_dim]) / window_strides[spatial_dim]) + 1.
• (C26) rank(result) = N.
• Nếu thao tác sử dụng các tensor không được định lượng:
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result).
• Nếu thao tác sử dụng các tensor được định lượng:
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C29) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) = kernel_output_feature_dimension.
  • (C30) Nếu is_per_axis_quantized(result), thì quantization_dimension(result) = output_feature_dimension.
  • Nếu is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs).
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C33) Nếu is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result).
  • Nếu !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result).

Ví dụ

// %lhs: [[
//        [
//          [1], [2], [5], [6]
//        ],
//        [
//          [3], [4], [7], [8]
//        ],
//        [
//          [10], [11], [14], [15]
//        ],
//        [
//          [12], [13], [16], [17]
//        ]
//      ]]
//
// %rhs: [
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]],
//        [[[1]], [[1]], [[1]]]
//       ]
%result = "stablehlo.convolution"(%lhs, %rhs) {
  window_strides = array<i64: 4, 4>,
  padding = dense<0> : tensor<2x2xi64>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  // In the StableHLO dialect, dimension numbers are encoded via:
  // `[<input dimensions>]x[<kernel dimensions>]->[output dimensions]`.
  // "b" is batch dimension, "f" is feature dimension,
  // "i" is input feature dimension, "o" is output feature dimension,
  // "0/1/etc" are spatial dimensions.
  dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
  batch_group_count = 1 : i64,
  feature_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[10], [26]],
//            [[46], [62]]
//          ]]

 Các ví dụ khác

cosin

Ngữ nghĩa

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %operand: [
//            [0.0, 1.57079632],       // [0, pi/2]
//            [3.14159265, 4.71238898] // [pi, 3pi/2]
//           ]
%result = "stablehlo.cosine"(%operand) : (tensor<2x2xf32>) -> tensor<2x2xf32>
// %result: [[1.0, 0.0], [-1.0, 0.0]]

 Các ví dụ khác

count_leading_zeros

Ngữ nghĩa

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

custom_call

Ngữ nghĩa

Đóng gói một thao tác do quá trình triển khai xác định call_target_name nhận inputs và called_computations, đồng thời tạo ra results. has_side_effect, backend_config và api_version có thể được dùng để cung cấp siêu dữ liệu bổ sung do quá trình 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á lộn xộn, phản ánh sự phát triển tự nhiên của thao tác tương ứng trong trình biên dịch XLA. Trong tương lai, chúng tôi dự định sẽ hợp nhất siêu dữ liệu này (#741).

Thông tin đầu vào

Kết quả đầu ra

(Hỗ trợ GPU XLA) Các mục tiêu custom_call đặc biệt

Có 3 call_target_name đặc biệt liên quan đến các loại buffer: CreateBuffer tạo một buffer chưa được khởi chạy, Pin tạo một buffer đã được khởi chạy và Unpin giải phóng một buffer và trả về nội dung của buffer.

%uninitialized_buffer = "stablehlo.custom_call"() {
  call_target_name = "CreateBuffer",
  api_version = 4 : i32,
} : () -> memref<4xf64>

%initialized_buffer = "stablehlo.custom_call"(%init_value) {
  call_target_name = "Pin",
  api_version = 4 : i32,
} : (tensor<4xf64>) -> memref<4xf64>

%dealloc_buffer = "stablehlo.custom_call"(%initialized_buffer) {
  call_target_name = "Unpin",
  api_version = 4 : i32,
} : (memref<4xf64>) -> tensor<4xf64>

Email đại diện

Một số hoạt động custom_call có thể yêu cầu một phần trong đầu ra và một phần trong các toán hạng để dùng chung bộ nhớ. Bạn có thể biểu thị điều này qua output_operand_aliases. Biểu diễn cặp bí danh bao gồm một danh sách các chỉ mục bộ dữ liệu đầu ra biểu thị phần đầu ra và một operand_index cùng với một danh sách các chỉ mục bộ dữ liệu toán hạng biểu thị phần toán hạng. Danh sách các chỉ mục bộ dữ liệu đầu ra hoặc toán hạng sẽ trống nếu loại tương ứng không phải là loại tuple và có thể có độ dài tuỳ ý đối với loại bộ dữ liệu lồng nhau tuỳ ý. Điều này tương tự như phần biểu thị biệt hiệu XLA.

Phần đầu ra và phần đầu vào trong một cặp bí danh phải có cùng loại. Đối với các hoạt động custom_call không phải là lệnh gọi đến CreateBuffer, Pin và Unpin, toán hạng buffer có thể xuất hiện tối đa trong một cặp bí danh và đầu ra buffer phải xuất hiện trong một cặp bí danh.

Ví dụ

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

%updated_buffer = "stablehlo.custom_call"(%buffer) {
  call_target_name = "Update",
  api_version = 4 : i32,
  output_operand_aliases = [
    #stablehlo.output_operand_alias<output_tuple_indices = [],
      operand_index = 0,
      operand_tuple_indices = []>]
} : (memref<4xf64>) -> memref<4xf64>

chia

Ngữ nghĩa

Thực hiện phép chia từng phần tử của các tensor số 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ố mà mọi phần phân số đều bị loại bỏ.
• Đối với số thực: division từ IEEE-754.
• Đối với số phức: phép chia số phức.
• Đối với các loại được lượng tử hoá:
  • dequantize_op_quantize(divide, lhs, rhs, type(result)).

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

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

 Các ví dụ khác

dot_general

Ngữ nghĩa

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

Cụ thể hơn, result[result_index] = dot_product, trong đó:

• lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions].
• rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions].
• result_batching_index + result_lhs_index + result_rhs_index = result_index trong đó size(result_batching_index) = size(lhs_batching_dimensions), size(result_lhs_index) = size(lhs_result_dimensions) và size(result_rhs_index) = size(rhs_result_dimensions).
• transposed_lhs = transpose(lhs, lhs_batching_dimensions + lhs_result_dimensions + lhs_contracting_dimensions).
• transposed_lhs_slice = slice(transposed_lhs, result_batching_index + result_lhs_index + [:, ..., :]).
• reshaped_lhs_slice = reshape(transposed_lhs_slice, dims(lhs, lhs_contracting_dimensions)).
• transposed_rhs = transpose(rhs, rhs_batching_dimensions + rhs_result_dimensions + rhs_contracting_dimensions).
• transposed_rhs_slice = slice(transposed_rhs, result_batching_index + result_rhs_index + [:, ..., :]).
• reshaped_rhs_slice = reshape(transposed_rhs_slice, dims(rhs, rhs_contracting_dimensions)).
• dot_product = reduce( inputs=[multiply(reshaped_lhs_slice, reshaped_rhs_slice)], init_values=[constant(0, element_type(result))], dimensions=range(size(lhs_contracting_dimensions)), body=lambda x, y: add(x, y)).

Đối với các loại được lượng tử hoá, hãy thực hiện dequantize_op_quantize( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs, type(result)).

Đối với các loại được lượng tử hoá kết hợp, hãy thực hiện hybrid_dequantize_then_op( lambda lhs, rhs: dot_general(lhs, rhs, lhs_batching_dimensions, rhs_batching_dimensions, lhs_contracting_dimensions, rhs_contracting_dimensions, precision_config), lhs, rhs).

precision_config kiểm soát sự đánh đổi giữa tốc độ và độ chính xác cho các phép tính trên các phần phụ trợ của bộ tăng tốc. Đây có thể là một trong những giá trị sau (hiện tại, ngữ nghĩa của các giá trị enum này chưa được chỉ định rõ ràng, nhưng chúng tôi dự định giải quyết vấn đề này trong #755):

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

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

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

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

Ví dụ về các thuộc tính DotAlgorithm:

// Inputs are casted to tf32, and then accumulated in f32:
{lhs_precision_type = tf32,
 rhs_precision_type = tf32,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = false}


// bf16_6x: each input is decomposed to 3 bf16 components, then 6 dot operations are done on those components, and the result is accumulated in f32.
{lhs_precision_type = bf16,
 rhs_precision_type = bf16,
 accumulation_type = f32,
 lhs_component_count = 3,
 rhs_component_count = 3,
 num_primitive_operations = 6,
 allow_imprecise_accumulation = false}


// Inputs are (casted to) f8e5m2, and we accumulate in f32, but for some steps we may accumulate in lower precision.
{lhs_precision_type = f8e5m2,
 rhs_precision_type = f8e5m2,
 accumulation_type = f32,
 lhs_component_count = 1,
 rhs_component_count = 1,
 num_primitive_operations = 1,
 allow_imprecise_accumulation = true}

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

• (C1) size(lhs_batching_dimensions) = size(rhs_batching_dimensions).
• (C2) size(lhs_contracting_dimensions) = size(rhs_contracting_dimensions).
• (C3) is_unique(lhs_batching_dimensions + lhs_contracting_dimensions).
• (C4) is_unique(rhs_batching_dimensions + rhs_contracting_dimensions).
• (C5) 0 <= lhs_batching_dimensions < rank(lhs).
• (C6) 0 <= lhs_contracting_dimensions < rank(lhs).
• (C7) 0 <= rhs_batching_dimensions < rank(rhs).
• (C8) 0 <= rhs_contracting_dimensions < rank(rhs).
• (C9) dim(lhs, lhs_batching_dimensions...) = dim(rhs, rhs_batching_dimensions...).
• (C10) dim(lhs, lhs_contracting_dimensions...) = dim(rhs, rhs_contracting_dimensions...).
• (C11) size(precision_config) = 2.
• (C12) shape(result) = dim(lhs, lhs_batching_dimensions) + dim(lhs, lhs_result_dimensions) + dim(rhs, rhs_result_dimensions).
• Nếu thao tác sử dụng các tensor không được định lượng:
  • (C13) element_type(lhs) = element_type(rhs).
• Nếu thao tác sử dụng các tensor được định lượng:
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs).
  • (C15) zero_points(rhs) = 0.
  • (C16) Nếu is_per_axis_quantized(rhs), thì quantization_dimension(rhs) không nằm trong rhs_contracting_dimensions.
  • Nếu is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs).
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result).
  • (C19) Nếu is_per_tensor_quantized(rhs), thì is_per_tensor_quantized(result).
  • Nếu !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result).
• Nếu !is_empty_algorithm(lhs_precision_type, rhs_precision_type, accumulation_type, lhs_component_count, rhs_component_count, num_primitive_operations allow_imprecise_accumulation):
  • (C21) precision_config... = DEFAULT.
  • (C22) 0 < lhs_component_count.
  • (C23) 0 < rhs_component_count.
  • (C24) 0 < num_primitive_operations.

Ví dụ

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

 Các ví dụ khác

dynamic_broadcast_in_dim

Ngữ nghĩa

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

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

Thông tin đầu vào

Kết quả đầu ra

Giới hạn

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

Ví dụ

// %operand: [
//            [1, 2, 3]
//           ]
%operand = stablehlo.constant dense<[[1, 2, 3]]> : tensor<1x3xi64>
%output_dimensions = stablehlo.constant dense<[2, 3, 2]> : tensor<3xi64>
%result = "stablehlo.dynamic_broadcast_in_dim"(%operand, %output_dimensions) {
  broadcast_dimensions = array<i64: 2, 1>,
  known_expanding_dimensions = array<i64: 0>,
  known_nonexpanding_dimensions = array<i64: 1>
} : (tensor<1x3xi64>, tensor<3xi64>) -> tensor<2x3x2xi64>
// %result: [
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ],
//            [
//             [1, 1],
//             [2, 2],
//             [3, 3]
//            ]
//          ]

 Các ví dụ khác

dynamic_conv

Ngữ nghĩa

Thao tác này có chức năng giống hệt với thao tác tích chập, nhưng khoảng đệm được chỉ định một cách linh hoạt thông qua padding.

Thông tin đầu vào