ข้อกำหนดของ StableHLO

StableHLO คือชุดการดำเนินการสำหรับการดำเนินงานระดับสูง (HLO) ในเครื่อง การเรียนรู้ (ML) StableHLO ทำงานเป็นเลเยอร์ความสามารถในการพกพาระหว่าง เฟรมเวิร์ก ML และคอมไพเลอร์ ML: เฟรมเวิร์ก ML ที่ผลิตโปรแกรม StableHLO สามารถใช้ร่วมกับคอมไพเลอร์ ML ที่ใช้โปรแกรม StableHLO ได้

เป้าหมายของเราคือการลดความซับซ้อนและเร่งการพัฒนา ML ด้วยการสร้าง ความสามารถในการทำงานร่วมกันระหว่างเฟรมเวิร์ก ML ต่างๆ (เช่น TensorFlow, JAX และ PyTorch) และคอมไพเลอร์ ML (เช่น XLA และ IREE) สุดท้ายนี้ เอกสารระบุข้อกำหนดสำหรับภาษาโปรแกรม StableHLO

ข้อกำหนดนี้ประกอบด้วยส่วนหลักๆ 3 ส่วน อย่างแรก ฟิลด์ ส่วนโปรแกรมจะอธิบายโครงสร้างของโปรแกรม StableHLO ซึ่งประกอบด้วยฟังก์ชัน StableHLO ซึ่งประกอบไปด้วยการทำงานของ StableHLO ภายในโครงสร้างนั้น ส่วน Ops จะระบุความหมายของ ของ Google แต่ละราย ส่วนการดำเนินการจะให้ความหมายแก่ผู้ใช้ทั้งหมด การดำเนินการเหล่านี้ร่วมกันภายในโปรแกรม สุดท้าย ส่วนเครื่องหมายจะกล่าวถึงเครื่องหมายที่ใช้ในส่วน

หากต้องการดูข้อมูลจำเพาะจาก StableHLO รุ่นก่อนหน้า ให้เปิดที่เก็บที่ ผลงานที่ติดแท็กที่สนใจ ตัวอย่างเช่น ข้อกำหนดเฉพาะ StableHLO v0.19.0 หากต้องการดูการเปลี่ยนแปลงที่เกิดขึ้นในระดับย่อยของ StableHLO แต่ละเวอร์ชัน โปรดดูที่ ให้ลงชื่อเข้าสู่ระบบ VhloDialect.td

โปรแกรม

Program ::= {Func}

โปรแกรม StableHLO ประกอบด้วยฟังก์ชัน StableHLO ที่กําหนดเอง ด้านล่างคือตัวอย่างโปรแกรมที่มีฟังก์ชัน @main ซึ่งมีอินพุต 3 แบบ (%image, %weights และ %bias) และ 1 เอาต์พุต เนื้อหาของฟังก์ชัน มี 6 Ops

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

ฟังก์ชัน

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

ฟังก์ชัน StableHLO (เรียกอีกอย่างว่าฟังก์ชันที่มีชื่อ) มี ตัวระบุ อินพุต/เอาต์พุต และเนื้อความ ในอนาคต เราวางแผนที่จะ แนะนำข้อมูลเมตาเพิ่มเติมสำหรับฟังก์ชันเพื่อความเข้ากันได้ที่ดีขึ้น ด้วย HLO (#425, #626, #740, #744)

รหัสระบุ

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

ตัวระบุ StableHLO คล้ายกับตัวระบุในหลายโปรแกรม ซึ่งมีลักษณะเฉพาะ 2 ประการ ได้แก่ 1) ตัวระบุทั้งหมดมีซิกเกิล ระบุตัวระบุประเภทต่างๆ 2) ตัวระบุค่าอาจเป็น เป็นตัวเลขทั้งหมดเพื่อลดความซับซ้อนในการสร้างโปรแกรม StableHLO

ประเภท

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

ประเภท StableHLO ได้รับการจัดหมวดหมู่เป็นประเภทค่า (หรือที่เรียกว่า ประเภทคลาสแรก) ซึ่งแสดงค่า StableHLO และประเภทที่ไม่ใช่ค่า ซึ่งอธิบายองค์ประกอบอื่นๆ ของโปรแกรม ประเภท HLO ที่เสถียรจะคล้ายกับประเภทใน ภาษาโปรแกรมหลายภาษา ที่มีความพิเศษคือ StableHLO ลักษณะเฉพาะโดเมนซึ่งส่งผลให้เกิดผลลัพธ์ที่ผิดปกติบางอย่าง (เช่น ประเภทสเกลาร์ ไม่ใช่ประเภทค่า)

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

ประเภท Tensor แสดงถึง Tensor ซึ่งก็คืออาร์เรย์แบบหลายมิติ พวกเขามี รูปร่าง และประเภทองค์ประกอบ โดยที่รูปร่างแสดงถึงค่าที่ไม่ติดลบหรือ ขนาดมิติข้อมูลที่ไม่รู้จักตามลำดับจากน้อยไปมากของขนาดนั้นๆ มิติข้อมูล (หรือเรียกอีกอย่างว่าแกน) ซึ่งมีหมายเลขตั้งแต่ 0 ถึง R-1 จำนวนมิติข้อมูล R เรียกว่าอันดับ ตัวอย่างเช่น tensor<2x3xf32> คือ ประเภท tensor ที่มีรูปร่าง 2x3 และประเภทองค์ประกอบ f32 มี 2 มิติ (หรืออีกนัยหนึ่งคือ 2 แกน) - มิติข้อมูลที่ 0 และมิติข้อมูลที่ 1 ซึ่งมีขนาด เท่ากับ 2 และ 3 อันดับคือ 2

รูปร่างอาจเป็นข้อมูลเพียงบางส่วนหรือไม่ทราบ (ไดนามิก) เช่น tensor<?x2xf64> ไม่รู้จักบางส่วนและไม่รู้จัก tensor<?x?xf64> ไดนามิก ขนาดของมิติข้อมูลจะแสดงโดยใช้ ? ไม่สามารถยกเลิกการจัดอันดับรูปร่าง

ในอนาคต เรากำลังวางแผนที่จะสำรวจประเภท Tensor เพิ่มเติมนอกเหนือจาก ขนาดมิติข้อมูลและประเภทองค์ประกอบ เช่น เพื่อรวมการออกแบบ (#629) และการขาดแคลน (#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

ประเภทองค์ประกอบที่กำหนดปริมาณแสดงค่าที่เป็นจำนวนเต็มของประเภทพื้นที่เก็บข้อมูลใน ช่วงจาก storage_min ถึง storage_max (รวม) ที่สอดคล้องกับ ค่าจุดลอยตัวของประเภทที่แสดง สำหรับค่าจำนวนเต็ม i ที่ระบุ ค่าจุดลอยตัว f ที่เกี่ยวข้องจะคำนวณเป็น f = (i - zero_point) * scale โดยที่เรียก scale และ zero_point พารามิเตอร์การวัดปริมาณ storage_min และ storage_max เป็นค่าที่ไม่บังคับ ในไวยากรณ์ แต่มีค่าเริ่มต้นเป็น min_value(storage_type) และ max_value(storage_type) ตามลำดับ ประเภทองค์ประกอบที่เป็นปริมาณจะมีค่า ข้อจำกัดต่อไปนี้

• (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) หากเป็น is_empty(quantization_dimension) ให้ size(scales) = 1
• (C11) 0 <= quantization_dimension

ขณะนี้ QuantizationScale เป็นค่าคงที่จุดลอยตัว แต่มีค่า มีความสนใจอย่างมากในสเกลแบบจำนวนเต็ม ซึ่งแสดงด้วยตัวคูณและ กะ เรามีแผนที่จะสํารวจเรื่องนี้ในอนาคตอันใกล้ (#1404)

กำลังสนทนาเกี่ยวกับความหมายของ QuantizationZeroPoint รวมถึงประเภท ค่า และดูว่ามีแค่หรือ อาจมีหลายจุดในประเภท Tensor ที่เล็กลง อิงตาม ผลของการสนทนานี้ ข้อกำหนดเกี่ยวกับจุดศูนย์อาจมีการเปลี่ยนแปลง ในอนาคต (#1405)

การสนทนาที่ดำเนินอยู่อีกรายการเกี่ยวข้องกับความหมายของ QuantizationStorageMin และ QuantizationStorageMax เพื่อระบุว่าควรมีข้อจำกัดใด กำหนดให้กับค่าเหล่านี้และค่าของ Tensor (#1406)

สุดท้าย เรากำลังวางแผนที่จะสำรวจการแสดงสเกลที่ไม่ทราบค่า 0 เช่นเดียวกับที่เราวางแผนจะสำรวจสิ่งที่ไม่มีข้อมูล ขนาด (#1407)

ประเภท tensor ที่กำหนดแสดงถึง Tensor ที่มีองค์ประกอบที่เล็กลง เหล่านี้ Tensor จะเหมือนกับ Tensor ปกติทุกประการ ยกเว้นตรงที่ มีประเภทองค์ประกอบที่เล็กลงแทนที่จะเป็นประเภทองค์ประกอบปกติ

ใน Tensor ที่เล็กลง การปรับเชิงปริมาณอาจเป็นแบบ ต่อ Tensor ได้ ซึ่งหมายความว่า 1 scale และ zero_point สำหรับ tensor ทั้งหมดหรืออาจเป็นแบบต่อแกน ซึ่งหมายถึงการมี scales และ zero_points หลายคู่ โดยมี 1 คู่ต่อ 1 ส่วน มิติข้อมูล quantization_dimension อย่างเป็นทางการมากขึ้นใน tensor t ด้วยการวัดขนาดต่อแกนจะมี dim(t, quantization_dimension) ส่วน จาก quantization_dimension: t[:, ..., 0, ..., :], t[:, ..., 1, ..., :] เป็นต้น องค์ประกอบทั้งหมดในชิ้นส่วนที่ i ใช้ scales[i] และ zero_points[i] เป็น พารามิเตอร์การวัดปริมาณ ประเภท Tensor ที่ปรับตามปริมาณมีสิ่งต่อไปนี้ ข้อจำกัด:

• สำหรับปริมาณต่อ Tensor:
  • ไม่มีข้อจำกัดเพิ่มเติม
• สำหรับการวัดปริมาณต่อแกน
  • (C12) quantization_dimension < rank(self)
  • (C13) dim(self, quantization_dimension) = size(scales)

TokenType ::= 'token'

ประเภทโทเค็นแสดงถึงโทเค็น ซึ่งก็คือค่าทึบแสงที่ผลิตและใช้ จากการดำเนินการบางอย่าง ใช้โทเค็นเพื่อกำหนดลำดับการดำเนินการในการดำเนินงาน ตามที่อธิบายไว้ในส่วนการดำเนินการ

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

ประเภทของสิ่งมีชีวิตเดียวกันหมายถึงบุคคลที่ 1 เท่านั้น (Tuple) กล่าวคือ มีรายการที่ซ้ำกัน Tuples เป็นมรดก ฟีเจอร์ที่มีไว้เฉพาะสำหรับความเข้ากันได้กับ HLO โดยใน HLO ใช้เพื่อแสดงอินพุตและเอาต์พุตตัวแปร ใน StableHLO จะมีอินพุตแปรผันและ เอาต์พุตได้รับการสนับสนุนโดยค่าเริ่มต้น และการใช้ Tuple เพียงอย่างเดียวใน StableHLO คือการ แสดงถึง HLO ABI ที่ครอบคลุม เช่น T, tuple<T> และ tuple<tuple<T>> อาจแตกต่างกันอย่างมากขึ้นอยู่กับ การใช้งานของคุณ เราวางแผนที่จะทำการเปลี่ยนแปลงกับ HLO ABI ในอนาคต ซึ่งอาจทำให้เรานำประเภท Tuple ออกจาก StableHLO (#598)

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

ประเภทองค์ประกอบแสดงองค์ประกอบของประเภท Tensor ต่างจากหลายๆ โปรแกรม ภาษาเหล่านี้ไม่ใช่คลาสแรกใน StableHLO ซึ่งหมายความว่า โปรแกรม StableHLO จะแสดงค่าของประเภทเหล่านี้โดยตรงไม่ได้ (ด้วยเหตุนี้ มันมีเอกลักษณ์ในการแสดงค่าสเกลาร์ของประเภท T ด้วย Tensor แบบ 0 มิติ ประเภท tensor<T>)

• ประเภทบูลีนแสดงถึงค่าบูลีน true และ false
• ประเภทจำนวนเต็มเป็นได้ทั้งลายเซ็น (si) หรือไม่ลงชื่อ (ui) และมี หนึ่งในความกว้างบิตที่รองรับ (2, 4, 8, 16, 32 หรือ 64) ประเภท siN ที่ลงนามแสดงค่าจำนวนเต็มตั้งแต่ -2^(N-1) ถึง 2^(N-1)-1 ประเภท uiN รวมและไม่มีเครื่องหมายแสดงค่าจำนวนเต็มตั้งแต่ 0 ถึง 2^N-1 เท่านั้น
• ประเภทจุดทศนิยมอาจเป็นอย่างใดอย่างหนึ่งต่อไปนี้
  • f8E4M3FN และ f8E5M2 ประเภทที่สัมพันธ์กับ การเข้ารหัส E4M3 และ E5M2 ในรูปแบบ FP8 ตามที่อธิบายไว้ใน รูปแบบ FP8 สำหรับการเรียนรู้เชิงลึก
  • ประเภท f8E4M3FNUZ และ f8E5M2FNUZ ที่สอดคล้องกับ E4M3 และ E5M2 รูปแบบ FP8 ตามที่อธิบายไว้ใน รูปแบบตัวเลข 8 บิตสำหรับโครงข่ายระบบประสาทเทียมระดับลึก
  • f8E4M3B11FNUZ ประเภทที่สอดคล้องกับการเข้ารหัส E4M3 ของรูปแบบ FP8 อธิบายไว้ใน การฝึกและการอนุมานแบบไฮบริดของจุดลอยตัว 8 บิต (HFP8) สำหรับเครือข่ายประสาทเทียมระดับลึก
  • bf16 ประเภทที่สอดคล้องกับรูปแบบ bfloat16 ที่อธิบายไว้ใน BFloat16: เคล็ดลับสู่ประสิทธิภาพสูงบน Cloud TPU
  • f16, f32 และ f64 ประเภทที่เกี่ยวข้องตามลำดับ binary16 ("ความแม่นยำครึ่งหนึ่ง"), binary32 ("ความแม่นยำครั้งเดียว") และ binary64 ("ความแม่นยำสองเท่า") ที่อธิบายใน มาตรฐาน IEEE 754
  • ประเภท tf32 สอดคล้องกับรูปแบบ TensorFloat32 และมีการรองรับที่จำกัดใน StableHLO
• ประเภทที่ซับซ้อนหมายถึงค่าที่ซับซ้อนซึ่งมีส่วนจริง และส่วนจินตภาพขององค์ประกอบประเภทเดียวกัน กลุ่มที่รองรับ ประเภทคือ complex<f32> (ทั้ง 2 ส่วนอยู่ในประเภท f32) และ complex<f64> (ทั้ง 2 ส่วนเป็นประเภท f64)

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

ประเภทฟังก์ชันจะแสดงทั้งฟังก์ชันที่มีชื่อและไม่ระบุตัวตน ซึ่งมี ประเภทอินพุต (รายการประเภททางด้านซ้ายของ ->) และประเภทเอาต์พุต (รายการประเภททางด้านขวาของ ->) ในหลายโปรแกรม ภาษา ประเภทฟังก์ชันเป็นคลาสแรก แต่ไม่ใช่ใน StableHLO

StringType ::= 'string'

ประเภทสตริงแสดงลำดับของไบต์ ต่างจากหลายๆ โปรแกรม ประเภทสตริงไม่ใช่คลาสแรกใน StableHLO และใช้ ระบุข้อมูลเมตาแบบคงที่สำหรับองค์ประกอบต่างๆ ของโปรแกรม

การดำเนินการ

การดำเนินการ StableHLO (เรียกอีกอย่างว่า ops) แสดงชุดที่ปิดแล้ว ของการดำเนินการระดับสูงในโมเดลแมชชีนเลิร์นนิง ดังที่กล่าวไว้ข้างต้น รูปแบบคำสั่ง StableHLO ได้รับแรงบันดาลใจมาจาก MLIR อย่างมาก ซึ่งไม่จำเป็นต้องเป็นเช่นนั้น เป็นทางเลือกที่หลักการยศาสตร์ แต่ก็เป็นวิธีที่ดีที่สุดสำหรับเป้าหมายของ StableHLO ทำให้เฟรมเวิร์ก ML และคอมไพเลอร์ ML ทำงานร่วมกันได้มากขึ้น

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

การดำเนินการ StableHLO (หรือเรียกว่าการดำเนินการ) มีชื่อ อินพุต/เอาต์พุตและลายเซ็น ชื่อประกอบด้วยคำนำหน้า stablehlo. และ ช่วยจำ ซึ่งระบุ Ops ที่สนับสนุนโดยไม่ซ้ำกัน ดูด้านล่างสำหรับ ซึ่งเป็นรายการที่ครอบคลุมการดำเนินการทั้งหมดที่รองรับ

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

การดำเนินการจะใช้อินพุตและสร้างเอาต์พุต อินพุตได้รับการจัดหมวดหมู่เป็น ค่าอินพุต (คำนวณระหว่างการดำเนินการ), ฟังก์ชันอินพุต (ที่มีให้ แบบคงที่เนื่องจากในฟังก์ชัน StableHLO ไม่ใช่ค่าที่ดีที่สุด) และ แอตทริบิวต์อินพุต (ระบุแบบคงที่ด้วย) ประเภทอินพุตและเอาต์พุต ที่ปฏิบัติการหนึ่งๆ บริโภคและผลิตขึ้นนั้นจะขึ้นอยู่กับความสามารถในการจำของตัวละครนั้นๆ ตัวอย่างเช่น add op จะใช้ค่าอินพุต 2 ค่าและสร้างค่าเอาต์พุต 1 ค่า เมื่อเปรียบเทียบกันแล้ว การดำเนินการ select_and_scatter จะใช้ค่าอินพุต 3 ค่า ฟังก์ชันอินพุต 2 รายการ และ 3 แอตทริบิวต์อินพุต

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

ฟังก์ชันอินพุต (หรือเรียกอีกอย่างว่าฟังก์ชันที่ไม่ระบุตัวตน) จะ คล้ายกับฟังก์ชันที่มีชื่อ เว้นแต่ว่า 1) ฟังก์ชันไม่มีตัวระบุ (ดังนั้น ชื่อ "ไม่ระบุชื่อ"), 2) จะไม่ประกาศประเภทเอาต์พุต (ประเภทเอาต์พุตคือ จะอนุมานได้จากการดำเนินการ return ภายในฟังก์ชัน)

ไวยากรณ์สำหรับฟังก์ชันอินพุตมีส่วนที่ไม่มีการใช้งานในปัจจุบัน (โปรดดู Unused เวอร์ชันที่ใช้งานจริงด้านบน) ซึ่งมีความเข้ากันได้กับ MLIR ใน MLIR มีแนวคิดที่กว้างขึ้นเกี่ยวกับ "ภูมิภาค" ซึ่งสามารถมี "การบล็อก" ได้หลายรายการ ของการดำเนินการที่เชื่อมโยงกัน บล็อกเหล่านี้มีรหัสที่สอดคล้องกับ เป็นเวอร์ชันที่ใช้งานจริง Unused เพื่อให้แยกความแตกต่างได้ StableHLO ไม่มีการดำเนินการข้าม ดังนั้นส่วนที่เกี่ยวข้องของไวยากรณ์ MLIR คือ ไม่ได้ใช้ (แต่ยังคงอยู่)

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

แอตทริบิวต์อินพุตมีชื่อและค่าซึ่งเป็นหนึ่งในแอตทริบิวต์ที่รองรับ ค่าคงที่นี้ เป็นวิธีหลักในการระบุข้อมูลเมตาแบบคงที่สำหรับรายการ จากองค์ประกอบเหล่านี้ ตัวอย่างเช่น การดำเนินการ concatenate ใช้แอตทริบิวต์ dimension เพื่อ ระบุมิติข้อมูลที่มีการต่อค่าอินพุต ในทำนองเดียวกัน การดำเนินการ slice ใช้แอตทริบิวต์หลายรายการ เช่น start_indices และ limit_indices เพื่อระบุขอบเขตที่ใช้ในการแบ่งค่าอินพุต

ปัจจุบันโปรแกรม StableHLO ในป่านี้บางครั้งอาจมีคุณสมบัติ ซึ่งไม่ได้อธิบายไว้ในเอกสารนี้ ในอนาคต เราวางแผนที่จะ ดูดซับแอตทริบิวต์เหล่านี้เข้าสู่กลุ่มกลยุทธ์ StableHLO หรือยับยั้งแอตทริบิวต์เหล่านี้ ปรากฏในโปรแกรม StableHLO ในระหว่างนี้ ลองดูรายการต่อไปนี้ ดังนี้

• layout (#629)
• mhlo.frontend_attributes (#628)
• mhlo.sharding (#619)
• output_operand_aliases (#740)
• ข้อมูลเมตาของตำแหน่ง (#594)

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

Op signature ประกอบด้วยประเภทของค่าอินพุตทั้งหมด (รายการประเภทใน ด้านซ้ายมือของ ->) และประเภทของค่าเอาต์พุตทั้งหมด (รายการ ประเภทที่อยู่ทางด้านขวาของ ->) กล่าวอย่างเจาะจงก็คือ ประเภทอินพุตคือ ซ้ำซ้อน และเอาต์พุตประเภท ก็มักจะซ้ำซ้อนกันเกือบทุกครั้ง (เนื่องจากสำหรับ การดำเนินการ StableHLO ส่วนใหญ่ สามารถอนุมานประเภทเอาต์พุตได้จากอินพุต) อย่างไรก็ตาม การดำเนินการ ลายเซ็นเป็นส่วนหนึ่งของไวยากรณ์ StableHLO อย่างตั้งใจเพื่อให้เข้ากันได้กับ MLIR

ด้านล่างนี้เป็นตัวอย่างการดำเนินการที่ฟีเจอร์การจำคือ select_and_scatter ใช้ 3 ค่าอินพุต (%operand, %source และ %init_value) ฟังก์ชันอินพุต 2 รายการ และแอตทริบิวต์อินพุต 3 รายการ (window_dimensions, window_strides และ padding) โปรดสังเกตว่าลายเซ็นของ Opin จะมีเฉพาะประเภทของค่าอินพุตเท่านั้นได้อย่างไร (แต่ไม่ใช่ประเภทของฟังก์ชันอินพุตและแอตทริบิวต์ที่มีให้ในบรรทัด)

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

ค่าคงที่

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

ค่าคงที่ HLO ประกอบด้วยลิเทอรัลและประเภทที่แสดงร่วมกัน ค่า StableHLO โดยทั่วไป ประเภทจะเป็นส่วนหนึ่งของไวยากรณ์แบบคงที่ ยกเว้น เมื่อไม่ชัดเจน (เช่น ค่าคงที่แบบบูลีนมีประเภท i1 อย่างชัดเจน ส่วนค่าคงที่จำนวนเต็มมีได้หลายประเภท)

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

ค่าคงที่บูลีนแสดงถึงค่าบูลีน true และ false บูลีน ค่าคงที่จะเป็นประเภท i1

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

ค่าคงที่จำนวนเต็มแสดงค่าจำนวนเต็มผ่านสตริงที่ใช้เลขทศนิยม หรือ ค่าฐานสิบหก ฐานอื่นๆ เช่น ไม่สนับสนุนฐานสองหรือฐานแปด ค่าคงที่จำนวนเต็มมีข้อจำกัดต่อไปนี้

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

ค่าคงที่จุดทศนิยมจะแสดงค่าที่จุดลอยตัวผ่านสตริงที่ ใช้ทศนิยมหรือสัญกรณ์วิทยาศาสตร์ นอกจากนี้ สัญกรณ์เลขฐานสิบหกสามารถ ใช้เพื่อระบุบิตพื้นฐานโดยตรงในรูปแบบจุดลอยตัวของ ประเภทที่เกี่ยวข้อง ค่าคงที่จุดลอยตัวมีข้อจำกัดต่อไปนี้

• (C1) ถ้ามีการใช้สัญกรณ์ที่ไม่ใช่เลขฐานสิบหก is_wellformed(float_literal, float_type)
• (C2) หากใช้ค่าฐานสิบหก size(hexadecimal_digits) = num_bits(float_type) / 4

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

ค่าคงที่เชิงซ้อนแสดงค่าเชิงซ้อนโดยใช้รายการของส่วนจริง (มาก่อน) และส่วนจินตภาพ (มาก่อน) ตัวอย่างเช่น (1.0, 0.0) : complex<f32> แสดงถึง 1.0 + 0.0i และ (0.0, 1.0) : complex<f32> แสดงถึง 0.0 + 1.0i ลำดับของมิติข้อมูลเหล่านี้ บางส่วนจะได้รับการจัดเก็บไว้ในหน่วยความจำมีการกำหนดการใช้งาน ค่าคงที่ที่ซับซ้อน มีข้อจำกัดต่อไปนี้

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

ค่าคงที่ของ Tensor แสดงค่า Tensor โดยใช้รายการที่ซ้อนกันที่ระบุผ่าน รูปแบบ NumPy เช่น dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32> แสดงค่า tensor ที่มีการแมปต่อไปนี้จากดัชนีไปยังองค์ประกอบ {0, 0} => 1, {0, 1} => 2, {0, 2} => 3, {1, 0} => 4, {1, 1} => 5 {1, 2} => 6 ลำดับการจัดเก็บองค์ประกอบเหล่านี้ในหน่วยความจำคือ ที่กำหนดโดยการติดตั้งใช้งาน ค่าคงที่ของ Tensor มีข้อจำกัดต่อไปนี้

• (C1) has_syntax(tensor_literal, element_type(tensor_type)) ซึ่งมีดังนี้
  • 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)) ซึ่งมีดังนี้
  • 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:])
  • หรือไม่เช่นนั้น false

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

ค่าคงที่ tensor ที่ปรับขยายตัวแล้วจะแสดงค่า tensor ที่แบ่งเป็นปริมาณโดยใช้ค่าเดียวกัน สัญกรณ์เป็นค่าคงที่ tensor พร้อมองค์ประกอบที่ระบุเป็นค่าคงที่ของ ประเภทพื้นที่เก็บข้อมูล ค่าคงที่ tensor ที่ปรับตามจำนวนมีข้อจํากัดต่อไปนี้

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

สัญพจน์สตริงประกอบด้วยไบต์ที่ระบุโดยใช้อักขระ ASCII และ Escape Sequence แต่สามารถเข้าใจการเข้ารหัสได้ ดังนั้นการตีความ จำนวนไบต์ได้รับการกำหนดตามการใช้งาน ลิเทอรัลสตริงมีประเภท string

การดำเนินการ

ท้อง

อรรถศาสตร์

ทำการดำเนินการ abs แบบองค์ประกอบบน Tensor ของ operand และสร้าง result Tensor ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับจำนวนเต็มที่มีเครื่องหมาย: โมดูลัสจำนวนเต็ม
• สำหรับแบบลอย: abs จาก IEEE-754
• สำหรับจำนวนเชิงซ้อน: โมดูลัสเชิงซ้อน
• สำหรับประเภทที่เล็กลง: dequantize_op_quantize(abs, operand, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) shape(result) = shape(operand)
• (C2) baseline_element_type(result) มีความหมายดังต่อไปนี้
  • complex_element_type(element_type(operand)) หากเป็น is_complex(operand)
  • จ่าย baseline_element_type(operand)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

เพิ่ม

อรรถศาสตร์

ทำการบวก Tensor 2 ตัว lhs และ rhs และสร้าง result Tensor ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับบูลีน: ตรรกะ OR
• สำหรับจำนวนเต็ม: การบวกจำนวนเต็ม
• สำหรับแบบลอย: addition จาก IEEE-754
• สำหรับจำนวนเชิงซ้อน: การบวกเชิงซ้อน
• สำหรับประเภทที่เล็กลง: dequantize_op_quantize(add, lhs, rhs, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

• หากการดำเนินการใช้ Tensor ที่ไม่ใช่เชิงปริมาณ
  • (C1) type(lhs) = type(rhs) = type(result)
• หากการดำเนินการใช้ Tensor ที่เล็กลง
  • (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) หากเป็น is_per_axis_quantized(lhs) ให้ quantization_dimension(lhs) = quantization_dimension(result)
  • (C7) หากเป็น is_per_axis_quantized(rhs) ให้ quantization_dimension(rhs) = quantization_dimension(result)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

after_all

อรรถศาสตร์

ตรวจสอบว่าได้ดำเนินการสร้าง inputs ก่อน การดำเนินการที่ขึ้นอยู่กับ result การดำเนินการดังกล่าวไม่ได้ทำอะไร มีไว้เพื่อสร้างทรัพยากร Dependency จาก result ถึง inputs เท่านั้น

อินพุต

เอาต์พุต

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

all_gather

อรรถศาสตร์

เชื่อมต่อค่าภายในกลุ่มกระบวนการแต่ละกลุ่มในตารางกระบวนการ StableHLO ของ Tensor จำนวน operands จากแต่ละกระบวนการตาม all_gather_dim และสร้าง Tensor จำนวน results รายการ

การดำเนินการนี้จะแยกตารางกริดประมวลผล StableHLO ออกเป็น process_groups ซึ่ง โดยมีคำจำกัดความดังนี้

• cross_replica(replica_groups) หากเป็น channel_id <= 0 and use_global_device_ids = false
• cross_replica_and_partition(replica_groups) หากเป็น channel_id > 0 and use_global_device_ids = false
• flattened_ids(replica_groups) หากเป็น channel_id > 0 and use_global_device_ids = true

หลังจากนั้น ภายในแต่ละ process_group:

• operands...@receiver = [operand@sender for sender in process_group] สำหรับทุกคน receiverในprocess_group
• results...@process = concatenate(operands...@process, all_gather_dim) สำหรับทุกคน processในprocess_group

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) 0 <= all_gather_dim < rank(operands...)
• (C2) is_unique(replica_groups)
• (C3) size(replica_groups) มีความหมายดังต่อไปนี้
  • num_replicas หากใช้ cross_replica
  • num_replicas หากใช้ cross_replica_and_partition
  • num_processes หากใช้ flattened_ids
• (C4) 0 <= replica_groups < size(replica_groups)
• (C5) หากเป็น use_global_device_ids = true ให้ channel_id > 0
• (C6) type(results...) = type(operands...) ยกเว้น
  • dim(results..., all_gather_dim) = dim(operands..., all_gather_dim) * dim(process_groups, 1)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

all_reduce

อรรถศาสตร์

ภายในกลุ่มกระบวนการแต่ละกลุ่มในตารางกระบวนการ StableHLO ให้ใช้การลด ฟังก์ชัน computation เป็นค่าของ tensor ของ operands จากแต่ละกระบวนการ และสร้าง results Tensor

การดำเนินการนี้จะแยกตารางกริดประมวลผล StableHLO ออกเป็น process_groups ซึ่ง โดยมีคำจำกัดความดังนี้

• cross_replica(replica_groups) หากเป็น channel_id <= 0 and use_global_device_ids = false
• cross_replica_and_partition(replica_groups) หากเป็น channel_id > 0 and use_global_device_ids = false
• flattened_ids(replica_groups) หากเป็น channel_id > 0 and use_global_device_ids = true

หลังจากนั้น ภายในแต่ละ process_group:

• results...@process[result_index] = exec(schedule) สำหรับแผนผังไบนารีบางทรี schedule โดยที่
  • exec(node) = computation(exec(node.left), exec(node.right))
  • exec(leaf) = leaf.value
• schedule คือแผนผังไบนารีที่กำหนดโดยการติดตั้งใช้งาน ซึ่งมีลำดับ การส่งผ่านคือ to_destination_type(operands...@process_group...[result_index], type(func_inputs(computation)[0]))

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) is_unique(replica_groups)
• (C2) size(replica_groups) มีความหมายดังต่อไปนี้
  • num_replicas หากใช้ cross_replica
  • num_replicas หากใช้ cross_replica_and_partition
  • num_processes หากใช้ flattened_ids
• (C3) 0 <= replica_groups < size(replica_groups)
• (C4) หากเป็น use_global_device_ids = true ให้ channel_id > 0
• (C5) computation มีประเภท (tensor<E>, tensor<E>) -> (tensor<E>) โดยที่ is_promotable(element_type(operand), E)
• (C6) shape(results...) = shape(operands...)
• (C7) element_type(results...) = E

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

all_to_all

อรรถศาสตร์

ภายในกลุ่มกระบวนการแต่ละกลุ่มในตารางการประมวลผล StableHLO ให้แยกค่าของ Tensor สี operands ตาม split_dimension แบ่งออกเป็นส่วนต่างๆ กระจายตัว เพื่อนำกระบวนการบางส่วนมาต่อกัน ส่วนที่กระจัดกระจายอยู่ตาม concat_dimension และสร้าง results Tensor การดำเนินการนี้จะแยกตารางกริดประมวลผล StableHLO ออกเป็น process_groups ซึ่ง โดยมีคำจำกัดความดังนี้

• cross_replica(replica_groups) หากเป็น channel_id <= 0
• cross_partition(replica_groups) หากเป็น channel_id > 0

หลังจากนั้น ภายในแต่ละ process_group:

• split_parts...@sender = split(operands...@sender, split_count, split_dimension) สำหรับ sender ทั้งหมดใน process_group
• scattered_parts...@receiver = [split_parts...@sender[receiver_index] for sender in process_group] ที่ receiver_index = process_group.index(receiver)
• results...@process = concatenate(scattered_parts...@process, concat_dimension)

อินพุต

เอาต์พุต

ข้อจำกัด

• (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) มีความหมายดังต่อไปนี้
  • num_replicas หากใช้ cross_replica
  • num_partitions หากใช้ cross_partition
• (C7) 0 <= replica_groups < size(replica_groups)
• (C8) dim(replica_groups, 1) = split_count
• (C9) type(results...) = type(operands...) ยกเว้นในกรณีที่ split_dimension != concat_dimension
  • dim(results..., split_dimension) = dim(operands..., split_dimension) / split_count
  • dim(results..., concat_dimension) = dim(operands..., concat_dimension) * split_count

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

และ

อรรถศาสตร์

ดำเนินการ AND ตามองค์ประกอบที่มี 2 Tensor คือ lhs และ rhs และสร้าง result Tensor ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับบูลีน: ตรรกะ AND
• สำหรับจำนวนเต็ม: Bitwise AND

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

อรรถศาสตร์

ดำเนินการ atan2 ตามองค์ประกอบใน lhs และ rhs tensor และสร้าง result Tensor ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับแบบลอย: atan2 จาก IEEE-754
• สำหรับจำนวนเชิงซ้อน: atan2 เชิงซ้อน
• สำหรับประเภทที่เล็กลง: dequantize_op_quantize(atan2, lhs, rhs, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

batch_norm_grad

อรรถศาสตร์

คำนวณการไล่ระดับสีของอินพุตหลายรายการของ Backpropagating ของ batch_norm_training จาก grad_output และผลิต grad_operand, grad_scale และ grad_offset Tensor อย่างเป็นทางการมากขึ้น การทำงานนี้สามารถแสดงเป็นการย่อยสลายเป็น การดำเนินการ StableHLO ที่มีอยู่โดยใช้ไวยากรณ์ Python ดังนี้

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

สำหรับประเภทที่เล็กลง จะดำเนินการ 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))

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) 0 <= feature_index < rank(operand)
• (C2) operand, scale, mean, variance, grad_output, grad_operand grad_scale และ grad_offset มี baseline_element_type เดียวกัน
• (C3) operand, grad_output และ grad_operand มีรูปร่างเดียวกัน
• (C4) scale, mean, variance, grad_scale และ grad_offset มี รูปทรงเดียวกัน
• (C5) size(scale) = dim(operand, feature_index)

ตัวอย่าง

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

อรรถศาสตร์

ปรับ tensor ของ operand ให้เป็นมาตรฐานในทุกมิติข้อมูลยกเว้น feature_index และสร้าง Tensor ขนาด result อย่างเป็นทางการ สามารถแสดงเป็นการแยกส่วนไปยังการดำเนินการ StableHLO ที่มีอยู่ โดยใช้ไวยากรณ์ Python ดังนี้

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)

สำหรับประเภทที่เล็กลง จะดำเนินการ 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))

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) 0 <= feature_index < rank(operand)
• (C2) operand, scale, offset, mean, variance และ result มี 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)

ตัวอย่าง

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

อรรถศาสตร์

คำนวณค่าเฉลี่ยและความแปรปรวนในมิติข้อมูลทั้งหมด ยกเว้น feature_index และปรับ tensor ของ operand ให้เป็นมาตรฐาน ซึ่งได้แก่ output, batch_mean และ batch_var Tensor อย่างเป็นทางการมากขึ้น การดำเนินงานนี้สามารถกำหนดเป็น ย่อยให้กับการดำเนินการ StableHLO ที่มีอยู่โดยใช้ไวยากรณ์ Python เป็น ดังต่อไปนี้:

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

สำหรับประเภทที่เล็กลง จะดำเนินการ 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))

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) 0 <= feature_index < rank(operand)
• (C2) operand, scale, offset, batch_mean, batch_var และ output 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)

ตัวอย่าง

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

อรรถศาสตร์

ดำเนินการบิตแคสต์บน tensor ของ operand และสร้าง Tensor ขนาด result ซึ่งบิตของ Tensor ทั้ง operand จะถูกตีความใหม่โดยใช้ ของ Tensor result

อย่างเป็นทางการมากขึ้น ให้ E = element_type(operand), E' = element_type(result) และ R = rank(operand):

• หากเป็น num_bits(E') < num_bits(E) bits(result[i0, ..., iR-1, :]) = bits(operand[i0, ..., iR-1])
• หากเป็น num_bits(E') > num_bits(E) bits(result[i0, ..., iR-2]) = bits(operand[i0, ..., iR-2, :])
• หากเป็น num_bits(E') = num_bits(E) bits(result[i0, ..., iR-1]) = bits(operand[i0, ..., iR-1])

bits แสดงผลการแทนค่าในหน่วยความจำของค่าที่ระบุและลักษณะการทำงานของค่าดังกล่าว มีการกำหนดการใช้งานเนื่องจากการนำเสนอ tensor ที่แน่นอนคือ การกำหนดการติดตั้ง และการนำเสนอประเภทองค์ประกอบที่แน่นอนคือ แบบกำหนดเองได้ด้วย

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) จาก E = is_quantized(operand) ? storage_type(operand) : element_type(operand), E' = is_quantized(result) ? storage_type(result) : element_type(result) และ R = rank(operand):
  • หากเป็น num_bits(E') = num_bits(E) shape(result) = shape(operand)
  • หาก num_bits(E') < num_bits(E):
  • rank(result) = R + 1
  • dim(result, i) = dim(operand, i) สำหรับ 0 <= i < R ทั้งหมด
  • dim(result, R) * num_bits(E') = num_bits(E)
  • หาก num_bits(E') > num_bits(E):
  • rank(result) = R - 1
  • dim(result, i) = dim(operand, i) สำหรับ 0 <= i < R ทั้งหมด
  • dim(operand, R - 1) * num_bits(E) = num_bits(E')
• (C2) หาก is_complex(operand) or is_complex(result) แล้ว is_complex(operand) and is_complex(result)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

broadcast_in_dim

อรรถศาสตร์

ขยายมิติข้อมูลและ/หรืออันดับของ Tensor อินพุตด้วยการทำซ้ำข้อมูล ใน tensor ของ operand และสร้าง Tensor ขนาด result อย่างเป็นทางการ result[result_index] = operand[operand_index] ตำแหน่งสำหรับ d ทั้งหมดใน axes(operand):

• operand_index[d] = 0 หากเป็น dim(operand, d) = 1
• จ่าย operand_index[d] = result_index[broadcast_dimensions[d]]

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) element_type(result) ได้มาจาก
  • element_type(operand) หากเป็น !is_per_axis_quantized(operand)
  • element_type(operand) ยกเว้น quantization_dimension(operand) ดังกล่าว scales(operand) และ zero_points(operand) อาจแตกต่างไปจาก quantization_dimension(result), scales(result) และ zero_points(result) ไม่เช่นนั้น
• (C2) size(broadcast_dimensions) = rank(operand)
• (C3) 0 <= broadcast_dimensions < rank(result)
• (C4) is_unique(broadcast_dimensions)
• (C5) สำหรับ d ทั้งหมดใน axes(operand):
  • dim(operand, d) = 1หรือ
  • dim(operand, d) = dim(result, broadcast_dimensions[d])
• (C6) หาก is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]
  • หากเป็น dim(operand, quantization_dimension(operand)) = 1 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

เคส

อรรถศาสตร์

สร้างเอาต์พุตจากการเรียกใช้ฟังก์ชันเดียวจาก branches ขึ้นอยู่กับค่าของ index อย่างเป็นทางการเพิ่มเติม result = selected_branch() โดยมี

• selected_branch = branches[index] หากเป็น 0 <= index < size(branches)
• จ่าย selected_branch = branches[-1]

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

CBT

อรรถศาสตร์

ดำเนินการรากลูกบาศก์ที่ชี้องค์ประกอบบน tensor ของ operand และสร้าง result Tensor ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับแบบลอย: rootn(x, 3) จาก IEEE-754
• สำหรับจำนวนเชิงซ้อน: รากที่สามเชิงซ้อน
• สำหรับประเภทที่เล็กลง: dequantize_op_quantize(cbrt, operand, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

Ceil

อรรถศาสตร์

สร้าง Ceil ที่ชี้ตามองค์ประกอบของ tensor ของ operand และสร้าง Tensor ขนาด result ใช้การดำเนินการ roundToIntegralTowardPositive จาก IEEE-754 สำหรับประเภทที่เล็กลง จะดำเนินการ dequantize_op_quantize(ceil, operand, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

Cholesky

อรรถศาสตร์

คำนวณการแตกตัวของ Cholesky ของกลุ่มเมทริกซ์

อย่างเป็นทางการมากขึ้นสำหรับ i ทั้งหมดใน index_space(result) result[i0, ..., iR-3, :, :] เป็นการแตกตัวของ Cholesky ของ a[i0, ..., iR-3, :, :] ในรูปของสามเหลี่ยมมุมล่างอันใดอันหนึ่ง (หาก lower คือ true) หรือเมทริกซ์สามเหลี่ยมด้านบน (หาก lower คือ false) ค่าเอาต์พุตในสามเหลี่ยมตรงข้าม เช่น สามเหลี่ยมมุมบนที่จำกัดหรือ สามเหลี่ยมมุมล่างแบบเข้มงวดที่สอดคล้องกันจะได้รับการกำหนดการใช้งาน

หากมี i ซึ่งเมทริกซ์อินพุตไม่ใช่ค่าจำกัดความบวกของเฮอร์มิเชียน เมทริกซ์ ลักษณะการทำงานจะไม่ได้ที่กำหนดไว้

สำหรับประเภทที่เล็กลง จะดำเนินการ dequantize_op_quantize(lambda operand: cholesky(operand, lower), a, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

แบบหนีบ

อรรถศาสตร์

ยึดทุกองค์ประกอบของ tensor ของ operand ระหว่างค่าต่ำสุดและสูงสุด และสร้าง Tensor เป็น result อย่างเป็นทางการเพิ่มเติมคือ result[result_index] = minimum(maximum(operand[result_index], min_element), max_element) โดยที่ min_element = rank(min) = 0 ? min[] : min[result_index] max_element = rank(max) = 0 ? max[] : max[result_index] สำหรับประเภทที่เล็กลง ทำการ dequantize_op_quantize(clamp, min, operand, max, type(result))

การใส่ลำดับด้วยจำนวนเชิงซ้อนนั้นเกี่ยวข้องกับอรรถศาสตร์ที่น่าประหลาดใจ ดังนั้นในอนาคต เราจึงวางแผนที่จะยกเลิกการสนับสนุนจำนวนเชิงซ้อน สำหรับการดำเนินการนี้ (#560)

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

collective_broadcast

อรรถศาสตร์

ภายในกลุ่มกระบวนการแต่ละกลุ่มในตารางการประมวลผล StableHLO ให้ส่งค่าของฟิลด์ operand Tensor ตั้งแต่กระบวนการต้นทางไปจนถึงกระบวนการเป้าหมายและสร้าง result Tensor

การดำเนินการนี้จะแยกตารางกริดประมวลผล StableHLO ออกเป็น process_groups ซึ่ง โดยมีคำจำกัดความดังนี้

• cross_replica(replica_groups) หากเป็น channel_id <= 0
• cross_partition(replica_groups) หากเป็น channel_id > 0

หลังจากนั้น result@process จะได้รับรายการโดย:

• operand@process_groups[i, 0] หากมี i ที่กระบวนการนี้เกิดขึ้น ใน process_groups[i]
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) หรือไม่เช่นนั้น

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) is_unique(replica_groups)
• (C2) 0 <= replica_groups < N โดยที่ N มีความหมายดังนี้
  • num_replicas หากใช้ cross_replica
  • num_partitions หากใช้ cross_partition
• (C3) type(result) = type(operand)

ตัวอย่าง

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

อรรถศาสตร์

ภายในกลุ่มกระบวนการแต่ละกลุ่มในตารางการประมวลผล StableHLO ให้ส่งค่าของแอตทริบิวต์ operand Tensor จากกระบวนการต้นทางไปยังกระบวนการเป้าหมายและสร้าง result Tensor

การดำเนินการนี้จะแยกตารางกริดประมวลผล StableHLO ออกเป็น process_groups ซึ่ง โดยมีคำจำกัดความดังนี้

• cross_replica(source_target_pairs) หากเป็น channel_id <= 0
• cross_partition(source_target_pairs) หากเป็น channel_id > 0

หลังจากนั้น result@process จะได้รับรายการโดย:

• operand@process_groups[i, 0] หากมี i ที่เป็นเช่นนั้น process_groups[i, 1] = process
• broadcast_in_dim(constant(is_quantized(result) ? quantize(0, element_type(result)) : 0, element_type(result)), [], type(result)) หรือไม่เช่นนั้น

อินพุต

เอาต์พุต

ข้อจำกัด

• (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 โดยที่ N มีความหมายดังนี้
  • num_replicas หากใช้ cross_replica
  • num_partitions หากใช้ cross_partition
• (C5) type(result) = type(operand)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

เปรียบเทียบ

อรรถศาสตร์

ดำเนินการเปรียบเทียบ tensor ของ lhs และ rhs ตามองค์ประกอบของ comparison_direction และ compare_type และสร้าง Tensor จำนวน result

ค่าของ comparison_direction และ compare_type มีค่าต่อไปนี้ อรรถศาสตร์:

สำหรับประเภทองค์ประกอบบูลีนและจำนวนเต็ม

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

สำหรับประเภทเอลิเมนต์ที่มีจุดลอยตัวที่มี compare_type = FLOAT การดำเนินการดังกล่าวจะนำมาใช้ การดำเนินการ IEEE-754 ต่อไปนี้

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

สำหรับประเภทเอลิเมนต์จุดลอยตัวที่มี compare_type = TOTALORDER การดำเนินการ ใช้ชุดค่าผสมของการดำเนินการ totalOrder และ compareQuietEqual จาก IEEE-754

สำหรับประเภทองค์ประกอบที่ซับซ้อน การเปรียบเทียบแบบพจนานุกรมของคู่ (real, imag) คือ ดำเนินการโดยใช้ comparison_direction และ compare_type ที่ระบุ การใส่ลำดับด้วยจำนวนเชิงซ้อนนั้นเกี่ยวข้องกับอรรถศาสตร์ที่น่าประหลาดใจ ดังนั้นในอนาคต เราจึงวางแผนที่จะยกเลิกการสนับสนุนจำนวนเชิงซ้อน เมื่อ comparison_direction เท่ากับ GE, GT, LE หรือ LT (#560)

สำหรับประเภทที่เล็กลง ทำการ dequantize_compare(lhs, rhs, comparison_direction)

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) baseline_element_type(lhs) = baseline_element_type(rhs)
• (C2) shape(lhs) = shape(rhs) = shape(result)
• (C3) compare_type มีความหมายดังต่อไปนี้
  • SIGNED หากเป็น is_signed_integer(element_type(lhs))
  • UNSIGNED หากเป็น is_unsigned_integer(element_type(lhs)) or is_boolean(element_type(lhs))
  • FLOAT หรือ TOTALORDER หากเป็น is_float(element_type(lhs))
  • FLOAT หากเป็น is_complex(element_type(lhs))

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

ซับซ้อน

อรรถศาสตร์

ทำ Conversion ตามองค์ประกอบเป็นค่าเชิงซ้อนจากคู่ของค่าจริงและ ค่าจินตภาพ lhs และ rhs แล้วสร้าง Tensor ขึ้น result

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) type(lhs) = type(rhs)
• (C2) shape(result) = shape(lhs)
• (C3) element_type(result) มีประเภท complex<E> โดยที่ E = element_type(lhs)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

การผสม

อรรถศาสตร์

สรุปการดำเนินการที่สร้างขึ้น (ประกอบขึ้น) ของการดำเนินการ StableHLO อื่นๆ การใช้ inputs และ composite_attributes และการผลิต results ตรรกะของทัศนะจะใช้โดยแอตทริบิวต์ decomposition สามารถแทนที่กระบวนการ composite ด้วยการแยกตัวออกโดยไม่ต้องเปลี่ยนโปรแกรม อรรถศาสตร์ ในกรณีที่ในบรรทัดการแตกตัวไม่ได้ให้ข้อมูล อรรถศาสตร์ทางเลือก ควรใช้ custom_call

ฟิลด์ version (ค่าเริ่มต้นคือ 0) ใช้เพื่อแสดงถึงค่าเมื่อองค์ประกอบ อรรถศาสตร์เปลี่ยนไป

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

เชื่อมต่อ

อรรถศาสตร์

เชื่อม inputs ตามมิติข้อมูล dimension ในลำดับเดียวกับ และสร้าง tensor ของ result อย่างเป็นทางการ result[i0, ..., id, ..., iR-1] = inputs[k][i0, ..., kd, ..., iR-1] โดยที่

1. id = d0 + ... + dk-1 + kd
2. d เท่ากับ dimension และ d0, ... เป็นขนาดมิติข้อมูลที่ d จาก inputs

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) same(element_type(inputs...))
• (C2) same(shape(inputs...)) ยกเว้น 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]) ยกเว้น
  • dim(result, dimension) = dim(inputs[0], dimension) + ...

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

ค่าคงที่

อรรถศาสตร์

สร้าง tensor ของ output จาก value คงที่

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

ทำให้เกิด Conversion

อรรถศาสตร์

ทำ Conversion ตามองค์ประกอบจากประเภทองค์ประกอบหนึ่งไปยังอีกประเภทหนึ่งใน operand Tensor และสร้าง Tensor ขนาด result

สำหรับ Conversion boolean-to-any-supported-type ค่า false จะเป็น แปลงเป็น 0 และแปลงค่า true เป็น 1 สำหรับ any-supported-type-to-boolean ระบบจะแปลงค่าเป็น 0 false และค่าที่ไม่ใช่ 0 จะแปลงเป็น true โปรดดูวิธีดำเนินการด้านล่าง ใช้งานได้สำหรับประเภทที่ซับซ้อน

สำหรับ Conversion ที่เกี่ยวข้องกับจำนวนเต็มถึงจำนวนเต็ม จำนวนเต็มถึงจุดทศนิยม หรือfloating-point-to-floating-point หากค่าของแหล่งที่มามีค่าเท่ากับ ที่แสดงในประเภทปลายทาง ค่าผลลัพธ์จะเท่ากับ การให้คำแนะนำ มิฉะนั้น ระบบจะแจ้งลักษณะการทำงานในภายหลัง (#180)

สำหรับ Conversion ที่เกี่ยวข้องกับfloating-point-to-integer ส่วนที่เป็นเศษส่วนจะเป็น ที่ถูกตัด หากไม่สามารถแสดงค่าที่ตัดในประเภทปลายทาง จะมีลักษณะการทำงานเดิมเมื่อใด (#180)

Conversion ที่เกี่ยวข้องกับซับซ้อนไปซับซ้อนมีลักษณะการทำงานเหมือนกับของ floating-point-to-floating-point สำหรับการแปลงค่าจริงและ ส่วนจินตภาพ

สำหรับ Conversion complex-to-any-other-type และ Conversion complex-to-any-other-type ระบบจะไม่สนใจค่าจินตภาพต้นทาง หรือค่าจินตภาพปลายทางคือ เป็นศูนย์ตามลำดับ การแปลงส่วนจริงเป็นไปตาม Conversion แบบจุดลอยตัว

โดยหลักการแล้ว การดำเนินการนี้อาจแสดงการลดปริมาณ (Conversion จาก Quenized tensor เป็น tensor ปกติ) การแปลงปริมาณ (Conversion จากปกติ tensor เป็น Tensor ที่เล็กลง) และการแปลงปริมาณอีกครั้ง (การแปลงระหว่างการวัดที่เล็กลง tensors) แต่ในขณะนี้เรามีการดำเนินการเฉพาะสำหรับเรื่องดังกล่าว uniform_dequantizeสำหรับ Use Case แรกและ uniform_quantize สำหรับ กรณีการใช้งานที่ 2 และ 3 ในอนาคต ทั้ง 2 ทางอาจมีการรวมเข้าด้วยกัน ลงใน convert (#1576)

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

คอนโวลูชัน

อรรถศาสตร์

ประมวลผลผลิตภัณฑ์แบบจุดระหว่างหน้าต่างของ lhs และส่วนย่อยของ rhs และสร้าง result แผนภาพต่อไปนี้แสดงวิธีการคำนวณองค์ประกอบใน result lhs และ rhs โดยใช้ตัวอย่างที่เป็นรูปธรรม

พิจารณารูปแบบอินพุตใหม่ต่อไปนี้ในรูปของ lhs เพื่อให้แสดงกรอบเวลาของ 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)

การจัดเฟรมใหม่ใช้ฟังก์ชันผู้ช่วยต่อไปนี้

• 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] โดยที่ j[d] = i[permutation[d]]

ถ้าเป็น feature_group_count = 1 และ batch_group_count = 1 สำหรับทั้งหมด output_spatial_index ใน index_space(dim(result, output_spatial_dimensions...)) result[result_shape(:, output_spatial_index, :)] = dot_product โดยที่

• 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]) ดูเหมือนว่าจะยังไม่มีการใช้งานฟีเจอร์นี้ เราจึงวางแผนที่จะนำฟีเจอร์นี้ออก (#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])

หาก 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)

หาก 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) วันที่

สำหรับประเภทที่เล็กลง จะดำเนินการ 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))

สำหรับประเภทที่เล็กลงของจำนวนแบบผสม จะดำเนินการ 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)

อินพุต

เอาต์พุต

ข้อจำกัด

• (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) จาก 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) จาก 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) จาก 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) มีความหมายดังต่อไปนี้
  • dim(lhs, input_batch_dimension) / batch_group_count หากเป็น result_dim = output_batch_dimension
  • dim(rhs, kernel_output_feature_dimension) หากเป็น result_dim = output_feature_dimension
  • num_windows หรือไม่เช่นนั้น ที่:
  • 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
• หากการดำเนินการใช้ Tensor ที่ไม่ใช่เชิงปริมาณ
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result)
• หากการดำเนินการใช้ Tensor ที่เล็กลง
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)
  • (C29) หาก is_per_axis_quantized(rhs) จากนั้นจ่าย quantization_dimension(rhs) = kernel_output_feature_dimension
  • (C30) หาก is_per_axis_quantized(result) quantization_dimension(result) = output_feature_dimension
  • หาก is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs)
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)
  • (C33) หาก is_per_tensor_quantized(rhs) is_per_tensor_quantized(result)
  • หาก !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

โคไซน์

อรรถศาสตร์

ดำเนินการโคไซน์ตามองค์ประกอบบน tensor ของ operand และสร้างพารามิเตอร์ result Tensor ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับแบบลอย: cos จาก IEEE-754
• สำหรับจำนวนเชิงซ้อน: โคไซน์เชิงซ้อน
• สำหรับประเภทที่เล็กลง: dequantize_op_quantize(cosine, operand, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

count_leading_zeros

อรรถศาสตร์

ดำเนินการนับตามองค์ประกอบของจำนวนบิตนำใน operand tensor และสร้าง tensor ของ result

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

custom_call

อรรถศาสตร์

สรุปการดำเนินการที่กำหนดโดยการใช้งาน call_target_name ซึ่งใช้เวลา inputs และ called_computations และสร้าง results has_side_effect, อาจใช้ backend_config และ api_version เพื่อให้ ข้อมูลเมตาที่กำหนดโดยการติดตั้งใช้งาน

ขณะนี้ การดำเนินการนี้มีคอลเล็กชันของ ข้อมูลเมตาที่สะท้อนให้เห็นถึงวิวัฒนาการที่เกิดขึ้นเองของการดำเนินการคู่กันใน คอมไพเลอร์ XLA เราวางแผนที่จะรวมข้อมูลเมตานี้เข้าด้วยกันในอนาคต (#741)

อินพุต

เอาต์พุต

ตัวอย่าง

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

หาร

อรรถศาสตร์

ทำการแบ่งตัวของตัวหาร lhs และตัวหาร rhs ตัวตั้งและตัวหาร สร้าง Tensor ของ result ดำเนินการต่อไปนี้โดยขึ้นอยู่กับประเภทองค์ประกอบ

• สำหรับจำนวนเต็ม: การหารจำนวนเต็มที่สร้างผลหารพีชคณิตด้วยค่าใดๆ ทิ้งส่วนที่เป็นเศษส่วนแล้ว
• สำหรับแบบลอย: division จาก IEEE-754
• สำหรับจำนวนเชิงซ้อน: การหารเชิงซ้อน
• สำหรับประเภทที่เล็กลง
  • dequantize_op_quantize(divide, lhs, rhs, type(result))

อินพุต

เอาต์พุต

ข้อจำกัด

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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

dot_general

อรรถศาสตร์

ประมวลผลผลิตภัณฑ์แบบจุดระหว่างส่วนต่างๆ ของ lhs และส่วนย่อยของ rhs และสร้าง tensor ของ result

อย่างเป็นทางการเพิ่มเติมคือ result[result_index] = dot_product ซึ่ง

• 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 โดยที่ 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)) วันที่

สำหรับประเภทที่เล็กลง จะดำเนินการ 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))

สำหรับประเภทที่เล็กลงของจำนวนแบบผสม จะดำเนินการ 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 จะควบคุมข้อดีข้อเสียระหว่างความเร็วและความแม่นยำสำหรับ การคำนวณในแบ็กเอนด์ของ Accelerator ซึ่งอาจเป็นอย่างใดอย่างหนึ่งต่อไปนี้ (ที่ แม้อรรถศาสตร์ของค่า enum เหล่านี้ก็ยังไม่ได้ระบุ แต่เราสามารถ มีแผนที่จะจัดการเรื่องนี้ใน #755):

• DEFAULT: การคำนวณที่เร็วที่สุด แต่ค่าประมาณที่แม่นยำที่สุดในช่วง หมายเลขเดิม
• HIGH: การคำนวณช้าลง แต่การประมาณการแม่นยำกว่า หมายเลขเดิม
• HIGHEST: การคำนวณช้าที่สุด แต่ประมาณที่แม่นยำที่สุดกับ หมายเลขเดิม

DotAlgorithm ระบุพร็อพเพอร์ตี้หลักของอัลกอริทึมที่จะนำมาใช้ จะใช้จุด ซึ่งจะกำหนดความแม่นยำด้วย หากแอตทริบิวต์อัลกอริทึม ได้รับการตั้งค่าแล้ว precision_config ต้องเป็น DEFAULT DotAlgorithms ไม่มีค่าเริ่มต้นเนื่องจากพารามิเตอร์เริ่มต้นคือการติดตั้งใช้งาน กำหนดไว้ จึงอาจตั้งค่าฟิลด์อัลกอริทึมจุดทั้งหมดเป็น None เพื่อระบุ อัลกอริทึมจุดว่าง ซึ่งจะใช้ค่า precision_config แทน

DotAlgorithm ช่องประกอบด้วย

• lhs_precision_type และ rhs_precision_type ความแม่นยำที่ LHS และ RHS ของการดำเนินการจะปัดเป็น ประเภทความแม่นยํานั้นไม่เกี่ยวข้องกับ ประเภทพื้นที่เก็บข้อมูลของอินพุตและเอาต์พุต
• accumulation_type ความแม่นยำที่ใช้ในการสะสม
• lhs_component_count rhs_component_count และnum_primitive_operations มีผลเมื่อเราสร้างอัลกอริทึมซึ่งจะสลาย LHS และ/หรือ RHS เป็น หลายคอมโพเนนต์ และดำเนินการ "พื้นฐาน" หลายอย่าง การดำเนินการกับจุดเหล่านั้น - มักเพื่อจำลองความแม่นยำที่สูงขึ้น (เช่น การใช้ประโยชน์จากประเภทข้อมูลปัญญาประดิษฐ์ (AI) bFloat 16 เพื่อการคำนวณที่แม่นยำยิ่งขึ้น: bf16_6x tf32_3x ฯลฯ) สำหรับอัลกอริทึมที่ไม่มีการแยกส่วน ค่าเหล่านี้ ควรตั้งค่าเป็น 1
• allow_imprecise_accumulation เพื่อระบุว่ามีการสะสมที่ความแม่นยำต่ำกว่าหรือไม่ ซึ่งทำได้ในบางขั้นตอน (เช่น CUBLASLT_MATMUL_DESC_FAST_ACCUM)

ตัวอย่างแอตทริบิวต์ 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}

ทั้งนี้ ขึ้นอยู่กับการติดตั้งใช้งานที่จะตัดสินใจว่าจะรองรับชุดค่าผสมใด ใน ไม่ได้รับประกันว่าอัลกอริทึมแต่ละรายการจะได้รับการสนับสนุนในแต่ละ ประเภท Accelerator ตามผู้บริโภค StableHLO หากอัลกอริทึมที่ระบุไม่ ข้อผิดพลาดควรเพิ่มขึ้น ไม่ใช่การกลับไปใช้ ทางเลือก การยืนยัน StableHLO จะให้การยืนยันที่ดีที่สุด ป้องกันอัลกอริทึมที่ยังไม่ได้รับการรองรับในฮาร์ดแวร์ใดๆ

โปรดดู xla_data.proto > Algorithm สำหรับค่าอัลกอริทึมที่รองรับบางค่า ตั๋ว #2483 แสดงแผนในการสร้าง เอกสารแบบรวมศูนย์บนอัลกอริทึมที่รองรับโดยแบ็กเอนด์

อินพุต

เอาต์พุต

ข้อจำกัด

• (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)
• หากการดำเนินการใช้ Tensor ที่ไม่ใช่เชิงปริมาณ
  • (C13) element_type(lhs) = element_type(rhs)
• หากการดำเนินการใช้ Tensor ที่เล็กลง
  • (C14) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)
  • (C15) zero_points(rhs) = 0
  • (C16) หาก is_per_axis_quantized(rhs) quantization_dimension(rhs) ไม่อยู่ใน rhs_contracting_dimensions
  • หาก is_quantized(lhs):
  • (C17) storage_type(lhs) = storage_type(rhs)
  • (C18) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)
  • (C19) หาก is_per_tensor_quantized(rhs) is_per_tensor_quantized(result)
  • หาก !is_quantized(lhs):
  • (C20) element_type(lhs) = expressed_type(rhs) = element_type(result)
• หาก !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

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

dynamic_broadcast_in_dim

อรรถศาสตร์

การดำเนินการนี้มีฟังก์ชันการทำงานเหมือนกับ broadcast_in_dim op แต่รูปร่างผลลัพธ์จะระบุแบบไดนามิกผ่าน output_dimensions

การดำเนินการนี้ยังยอมรับแอตทริบิวต์ที่ไม่บังคับ known_expanding_dimensions, known_non_expanding_dimensions ด้วย เพื่อแสดงความรู้แบบคงที่เกี่ยวกับพฤติกรรมการขยายมิติข้อมูล หากไม่ได้ระบุ ระบบจะถือว่ามิติข้อมูลทั้งหมดกำลังขยายได้

อินพุต

เอาต์พุต

ข้อจำกัด

• (C1) element_type(result) ได้มาจาก
  • element_type(operand) หากเป็น !is_per_axis_quantized(operand)
  • element_type(operand) ยกเว้น quantization_dimension(operand) ดังกล่าว scales(operand) และ zero_points(operand) อาจแตกต่างไปจาก quantization_dimension(result), scales(result) และ zero_points(result) ไม่เช่นนั้น
• (C2) size(broadcast_dimensions) = rank(operand)
• (C3) 0 <= broadcast_dimensions < rank(result)
• (C4) is_unique(broadcast_dimensions)
• (C5) สำหรับ d ทั้งหมดใน axes(operand):
  • dim(operand, d) = 1หรือ
  • dim(operand, d) = dim(result, broadcast_dimensions[d])
• (C6) หาก is_per_axis_quantized(result):
  • quantization_dimension(result) = broadcast_dimensions[quantization_dimension(operand)]
  • หากเป็น dim(operand, quantization_dimension(operand)) = 1 scales(result)[i] = scales(operand)[0] and zero_points(result)[i] = zero_points(operand)[0] for i in range(dim(result, quantization_dimension(result)))
• (C7) size(output_dimensions) = rank(result)
• (C8) is_unique(known_expanding_dimensions + known_non_expanding_dimensions)
• (C9) 0 <= known_expanding_dimensions < rank(operand)
• (C10) 0 <= known_non_expanding_dimensions < rank(operand)

ตัวอย่าง

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

ตัวอย่างเพิ่มเติม

dynamic_conv

อรรถศาสตร์

การดำเนินการนี้มีฟังก์ชันการทำงานเหมือนกับ อ๊ะ แต่ระบุระยะห่างจากขอบแบบไดนามิกผ่าน padding

อินพุต

เอาต์พุต

ข้อจำกัด

• (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) จาก 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) จาก 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) จาก 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) มีความหมายดังต่อไปนี้
  • dim(lhs, input_batch_dimension) / batch_group_count หากเป็น result_dim = output_batch_dimension
  • dim(rhs, kernel_output_feature_dimension) หากเป็น result_dim = output_feature_dimension
  • num_windows หรือไม่เช่นนั้น ที่:
  • 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
• หากการดำเนินการใช้ Tensor ที่ไม่ใช่เชิงปริมาณ
  • (C27) element_type(lhs) = element_type(rhs) = element_type(result)
• หากการดำเนินการใช้ Tensor ที่เล็กลง
  • (C28) is_quantized(lhs) = is_quantized(result) and is_quantized(rhs)
  • (C29) หาก is_per_axis_quantized(rhs) จากนั้นจ่าย quantization_dimension(rhs) = kernel_output_feature_dimension
  • (C30) หาก is_per_axis_quantized(result) quantization_dimension(result) = output_feature_dimension
  • หาก is_quantized(lhs):
  • (C31) storage_type(lhs) = storage_type(rhs)
  • (C32) expressed_type(lhs) = expressed_type(rhs) = expressed_type(result)
  • (C33) หาก is_per_tensor_quantized(rhs) is_per_tensor_quantized(result)
  • หาก !is_quantized(lhs):
  • (C34) element_type(lhs) = expressed_type(rhs) = element_type(result)

ตัวอย่าง

// %lhs: [[
//        [[1], [2], [5], [6]],
//        [[3], [4], [7], [8]],
//        [[10], [11], [14], [15]],
//        [[12], [13], [16], [17]]
//      ]]
//
// %rhs: [
//         [[[1]], [[1]], [[1]]],
//         [[[1]], [[1]], [[1]]],
//         [[[1]], [[1]], [[1]]]
//        ]
// %padding: [[1, 1],
//            [1, 1]]
%result = "stablehlo.dynamic_conv"(%lhs, %rhs, %padding) {
  window_strides = array<i64: 4, 4>,
  lhs_dilation = array<i64: 2, 2>,
  rhs_dilation = array<i64: 1, 1>,
  window_reversal = array<i1: false, false>,
  dimension_numbers = #stablehlo.conv<raw
    input_batch_dimension = 0,
    input_feature_dimension = 3,
    input_spatial_dimensions = [0, 1],
    kernel_input_feature_dimension = 2,
    kernel_output_feature_dimension = 3,
    kernel_spatial_dimensions = [0, 1],
    output_batch_dimension = 0,
    output_feature_dimension = 3,
    output_spatial_dimensions = [1, 2]
  >,
  feature_group_count = 1 : i64,
  batch_group_count = 1 : i64,
  precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>]
} : (tensor<1x4x4x1xi64>, tensor<3x3x1x1xi64>, tensor<2x2xi64>) -> tensor<1x2x2x1xi64>
// %result: [[
//            [[1], [5]],
//            [[10], [14]]
//          ]]