'sdy' ภาษาถิ่น

รูปแบบ Shardy (SDY) จะกำหนดการนำเสนอการแยกกลุ่มเทนเซอร์ตามแกน และคอมโพเนนต์ API เพิ่มเติมเพื่อแนบการแยกกลุ่มกับเทนเซอร์

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

sdy.constant (sdy::ConstantOp)

การดำเนินการแบบคงที่

สร้างเทนเซอร์ output จากค่าคงที่ value

ดูที่ https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

ตัวอย่าง

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

ลักษณะ: AlwaysSpeculatableImplTrait

อินเทอร์เฟซ: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

ผลลัพธ์: MemoryEffects::Effect{}

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
value::mlir::ElementsAttrแอตทริบิวต์เวกเตอร์/เทนเซอร์แบบคงที่

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
output เทนเซอร์ของค่าประเภทใดก็ได้

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

การดำเนินการที่ขอบของการรับส่งข้อมูล

ไวยากรณ์:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

ขอบการไหลของข้อมูลของการดำเนินการ X บางรายการจะกำหนดบริดจ์ระหว่างชุดแหล่งที่มา (แต่ละรายการเป็นโอเปอเรนด์ของ X หรือโอเปอเรนด์ของเงื่อนไขสิ้นสุดบล็อกของ X) และชุดเป้าหมาย (แต่ละรายการเป็นผลลัพธ์ของ X หรืออาร์กิวเมนต์บล็อกของ X) เพื่อให้มีการแบ่งกลุ่มแหล่งที่มาและเป้าหมายทั้งหมดในลักษณะเดียวกัน

การดำเนินการหนึ่งๆ อาจมีขอบการไหลของข้อมูลหลายรายการที่ตัดกัน

เช่น

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

การดำเนินการ while นี้มีขอบการรับส่งข้อมูล n รายการ โดยขอบการรับส่งข้อมูลลำดับที่ i อยู่ระหว่างแหล่งที่มา x_i, return_value_i และปลายทาง y_i, pred_arg_i, body_arg_i

sdy.data_flow_edge ใช้เป็นอินพุตเป้าหมายรูทของ Edge (อาจเป็นเป้าหมายใดก็ได้ แต่ควรเป็นผลลัพธ์ op แทนที่จะเป็นอาร์กิวเมนต์บล็อก) ซึ่งไม่ควรมีประโยชน์อื่นใด การดำเนินการนี้ไม่ใช่การดำเนินการแบบบริสุทธิ์เนื่องจากใช้อินพุตที่เดิมไม่มีการใช้งาน

sdy.data_flow_edge ยังมีการจัดสรรที่ไม่บังคับสําหรับเป้าหมายทั้งหมดของขอบด้วย และควรอัปเดตการจัดสรรนั้นแทนการจัดสรรของเป้าหมาย (หากแนบได้) ในระหว่างการนำไปใช้งาน วิธีนี้มีประโยชน์เมื่อโอเปอเรเตอร์ มีขอบมากมาย เนื่องจากจะให้ประสิทธิภาพดีกว่ามาก

  • จะแพร่กระจายผ่านขอบแต่ละด้านแยกกัน
  • อัปเดตการแยกแต่ละขอบแยกกันแทนที่จะอัปเดตเป้าหมายทั้งหมดพร้อมกัน (เช่น การดำเนินการมี TensorShardingPerValueAttr รายการเดียวที่แก้ไขไม่ได้สำหรับการแยกผลลัพธ์)
  • เพิ่มแต่ละขอบลงในเวิร์กลิสต์แยกกันเมื่อมีการแยกข้อมูลของแหล่งที่มา

การนำไปใช้งานจะนำไปใช้งานการแยกส่วนระหว่างแหล่งที่มาและเป้าหมายทั้งหมดของ sdy.data_flow_edge ราวกับว่าเป็นการดำเนินการปกติที่มีแหล่งที่มาเป็นตัวดำเนินการและเป้าหมายเป็นผลลัพธ์ และ sdy.op_sharding_rule ที่เป็นข้อมูลระบุตัวตน ซึ่งหมายความว่าการนำไปข้างหน้าจะมาจากแหล่งที่มาไปยังเป้าหมาย และการนำไปข้างหลังจะมาจากเป้าหมายไปยังแหล่งที่มา

เราไม่อนุญาตให้การป้อน sdy.data_flow_edge ได้รับการกําหนดโดย SdyDialect op เราจึงจะถือว่าการป้อนนั้นได้รับการกําหนดโดย op ที่แอตทริบิวต์ sdy.sharding ไม่ได้ลงทะเบียน

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: InferTypeOpInterface

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
sharding::mlir::sdy::TensorShardingAttrการแยกกลุ่ม Tensor

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
input รูปแบบของค่าประเภทใดก็ได้

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
result รูปแบบของค่าประเภทใดก็ได้

sdy.manual_computation (sdy::ManualComputationOp)

การดำเนินการแบบขนานในอุปกรณ์หลายเครื่องด้วยกลุ่มที่รวบรวมด้วยตนเอง

ไวยากรณ์:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

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

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

ลักษณะ: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
in_shardings::mlir::sdy::TensorShardingPerValueAttrการแยกกลุ่ม Tensor ตามออปเรอเรนด์/ผลลัพธ์ของการดำเนินการ
out_shardings::mlir::sdy::TensorShardingPerValueAttrการแยกกลุ่ม Tensor ตามออปเรอเรนด์/ผลลัพธ์ของการดำเนินการ
manual_axes::mlir::sdy::ManualAxesAttr

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
tensors ตัวแปรของเทนเซอร์การจัดอันดับของค่าประเภทใดก็ได้

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
results ตัวแปรของ tensor ที่มีการจัดอันดับของค่าทุกประเภท

sdy.mesh (sdy::MeshOp)

เมชที่มีชื่อ

ไวยากรณ์:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

กำหนดเมชที่มีชื่อใหม่ เมชทั้งหมดในโมดูลต้องมีจํานวนอุปกรณ์เท่ากัน (ยกเว้นเมชที่มี device_id รายการเดียว) เมชคือการดำเนินการ Symbol ที่ปรากฏใน SymbolTable ของข้อบังคับ และสามารถอ้างอิงด้วย name ของข้อบังคับ

ลักษณะ: HasParent<ModuleOp>

อินเทอร์เฟซ: Symbol

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
sym_name::mlir::StringAttrแอตทริบิวต์สตริง
mesh::mlir::sdy::MeshAttrตาข่ายของแกนและรายการอุปกรณ์

sdy.named_computation (sdy::NamedComputationOp)

การดำเนินการคำนวณที่มีชื่อ

ไวยากรณ์:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

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

ซึ่งสามารถใช้เพื่อจัดการการเผยแพร่ผ่านคำสั่งการเรียกไปยังฟังก์ชันอื่นๆ ผู้ใช้ Shardy ควรเขียนพาสการนําเข้า/ส่งออกที่จะแปลงการดําเนินการการเรียกไปยังการดําเนินการ sdy.named_computation โดยทำซ้ำ/คัดลอกเนื้อหาของฟังก์ชันที่เรียกไปยังเนื้อหาของ named_computation

ประเภทของอาร์กิวเมนต์บล็อกแต่ละรายการและค่าที่แสดงผลในภูมิภาคต้องเหมือนกับประเภทของออพเพอร์แลนด์และประเภทผลลัพธ์ของการดำเนินการ

ตัวอย่าง

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

ลักษณะ: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

อินเทอร์เฟซ: ConditionallySpeculatable, ShardableDataFlowOpInterface

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
name::mlir::StringAttrแอตทริบิวต์สตริง
in_shardings::mlir::sdy::TensorShardingPerValueAttrการแยกกลุ่ม Tensor ตามออปเรอเรนด์/ผลลัพธ์ของการดำเนินการ
out_shardings::mlir::sdy::TensorShardingPerValueAttrการแยกกลุ่ม Tensor ตามออปเรอเรนด์/ผลลัพธ์ของการดำเนินการ

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
operands แปรผันทุกประเภท

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
«ไม่มีชื่อ» ตัวแปรใดก็ได้

sdy.propagation_barrier (sdy::PropagationBarrierOp)

การดำเนินการกับสิ่งกีดขวางการนำไปใช้งาน

ไวยากรณ์:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

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

ซึ่งจะช่วยป้องกันไม่ให้มีการนำไปใช้การแยกกลุ่มระหว่างการใช้ผลลัพธ์ของการดำเนินการแบเรียร์กับโอเปอเรนด์

  • FORWARD หมายความว่าการแยกส่วนสามารถส่งผ่านจากออปเรอนด์ไปยังผลลัพธ์ได้เท่านั้น
  • BACKWARD หมายความว่าการแยกข้อมูลจะไหลจากผลลัพธ์ไปยังออบเจ็กต์ได้เท่านั้น
  • NONE หมายความว่าระบบไม่สามารถกระจายข้อมูลการแยกส่วนผ่านการดำเนินการนี้
  • ไม่สามารถระบุ BOTH ได้ เนื่องจากการดำเนินการนี้จะซ้ำซ้อน

ลักษณะ: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

อินเทอร์เฟซ: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

ผลลัพธ์: MemoryEffects::Effect{}

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
allowed_direction::mlir::sdy::PropagationDirectionAttrรายการทิศทางการนำไปใช้งาน

ตัวถูกดำเนินการ:

ตัวถูกดำเนินการ คำอธิบาย
input เทนเซอร์ที่มีการจัดอันดับของค่าประเภทใดก็ได้

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
result เทนเซอร์ที่มีการจัดอันดับของค่าประเภทใดก็ได้

sdy.reshard (sdy::ReshardOp)

แยกกลุ่มเทนเซอร์ใหม่เป็นกลุ่มอื่น

ไวยากรณ์:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

แบ่งกลุ่มเทนเซอร์อินพุตใหม่โดยใช้การแยกกลุ่มที่ระบุ ซึ่งแตกต่างจากการแยกกลุ่มที่มีอยู่ของเทนเซอร์อินพุต

ทั้ง ShardingConstraintOp และ ReshardOp จะแนบชาร์ดกับ Tensor ช่วงอายุของสัตว์เหล่านี้

  1. ก่อนการชาร์ดการเผยแพร่ ผู้ใช้จะเพิ่ม ShardingConstraintOp
  2. การนำไปใช้งานการแยกกลุ่มจะใช้ ShardingConstraintOp ไม่มี ShardingConstraintOp ในผลลัพธ์ของการนำไปใช้การแยกส่วน แต่ระบบอาจเพิ่ม ReshardOp หากจําเป็น
  3. พาร์ติชันแปลง ReshardOp เป็นกลุ่มรวม (หรือตัวเลือกอัตลักษณ์) ไม่ควรมี ReshardOp ในผลลัพธ์ของโปรแกรมแบ่งพาร์ติชัน

// TODO(b/331680067). เพิ่มรูปแบบการจัดทําให้เป็นรูปแบบมาตรฐานเพื่อนําการดําเนินการการจัดสรรใหม่ซ้ำซ้อนออก //

ลักษณะ: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

อินเทอร์เฟซ: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

ผลลัพธ์: MemoryEffects::Effect{}

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
sharding::mlir::sdy::TensorShardingAttrการแยกกลุ่ม Tensor

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
input เทนเซอร์ของค่าประเภทใดก็ได้

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
result เทนเซอร์ของค่าประเภทใดก็ได้

sdy.return (sdy::ReturnOp)

การดำเนินการ sdy.return จะสิ้นสุดการทำงานของภูมิภาคที่แนบอยู่กับการดำเนินการตามภูมิภาค sdy และการดำเนินการตามภูมิภาคแบบ Shardy อื่นๆ ฟังก์ชันนี้รับอาร์กิวเมนต์เป็นลิสต์ค่าที่มีประเภทใดก็ได้ (แต่ต้องเป็นแบบเดียวกัน เช่น AnyTensor) จึงนํามาใช้ซ้ำได้ในระดับต่างๆ ของสแต็ก IR แบบ Shardy

ไวยากรณ์:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

ลักษณะเฉพาะ: AlwaysSpeculatableImplTrait, Terminator

อินเทอร์เฟซ: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

ผลลัพธ์: MemoryEffects::Effect{}

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
results ตัวแปรใดก็ได้

sdy.sharding_constraint (sdy::ShardingConstraintOp)

จำกัดเทนเซอร์ให้อยู่ในการแยกข้อมูลที่กำหนด

ไวยากรณ์:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

แนบชาร์ดดิ้งกับ tensor ระดับกลาง (เช่น ผลลัพธ์ของ Matmul) เพื่อระบุว่าควรชาร์ด Tensor หรือบางส่วนของการใช้งานดังกล่าว

หากการแยกข้อมูลมีหลายมิติข้อมูลและแกนที่ไม่มีการจํากัด หมายความว่าเทนเซอร์จะแยกตามมิติข้อมูลแบบเปิดได้

การดำเนินการนี้จะดำเนินการอย่างใดอย่างหนึ่งต่อไปนี้

  • ไม่มีการใช้งาน (อันตราย) ซึ่งหมายความว่าการชาร์ดดิ้งที่แนบมาคือวิธีการชาร์ด Tensor อินพุต
  • มีการใช้งาน - ซึ่งหมายความว่าการแยกส่วนงานที่แนบมาคือวิธีที่ควรแยกส่วนของการใช้งานการดําเนินการข้อจํากัดการแยกส่วน ขณะที่การใช้งานอื่นๆ ของเทมพอร์ลอินพุตอาจมีการแยกส่วนที่แตกต่างกัน (หากเทมพอร์ลอินพุตไม่มีการใช้งานอื่นๆ ลักษณะการทํางานจะเหมือนกับกรณีไม่มีการใช้งาน)

ลักษณะ: Elementwise, SameOperandsAndResultType

อินเทอร์เฟซ: InferTypeOpInterface

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
sharding::mlir::sdy::TensorShardingAttrการแยกกลุ่ม Tensor

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
input เทนเซอร์ของค่าประเภทใดก็ได้

ผลลัพธ์:

ผลลัพธ์ คำอธิบาย
result เทนเซอร์ของค่าประเภทใดก็ได้

sdy.sharding_group (sdy::ShardingGroupOp)

การดำเนินการชาร์ดดิ้งกลุ่ม

ไวยากรณ์:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

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

แอตทริบิวต์

แอตทริบิวต์ประเภท MLIRคำอธิบาย
group_id::mlir::IntegerAttrแอตทริบิวต์จำนวนเต็มแบบไม่มีเครื่องหมาย 64 บิต

ออบเจ็กต์:

ตัวถูกดำเนินการ คำอธิบาย
input เทนเซอร์ที่มีการจัดอันดับของค่าประเภทใดก็ได้

Attributes

AxisRefAttr

การอ้างอิงถึงแกนทั้งแกนหรือแกนย่อยที่แยก

ไวยากรณ์:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
ชื่อ ::llvm::StringRef ชื่อ
sub_axis_info SubAxisInfoAttr

DimMappingAttr

รายการดัชนีปัจจัยสําหรับมิติข้อมูล

ดัชนีปัจจัยทั้งหมดต้องอยู่ในช่วง [0, num_factors) และรายการว่างเปล่าบ่งบอกว่านี่เป็นการแมปที่เป็นค่าว่าง (แยกวิเคราะห์/พิมพ์ด้วย *) หมายความว่ามิติข้อมูลไม่ได้จับคู่กับปัจจัยใดๆ

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
factor_indices ::llvm::ArrayRef<int64_t>

DimensionShardingAttr

ชาร์ดดิ้งมิติข้อมูล

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

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
แกน ::llvm::ArrayRef<AxisRefAttr> รายการการอ้างอิงแกน
is_closed bool
ลำดับความสำคัญ std::optional<int64_t>

ManualAxesAttr

ไวยากรณ์:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
value ::llvm::ArrayRef<StringAttr>

MeshAttr

ตาข่ายของแกนและรายการอุปกรณ์

ไวยากรณ์:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

เมชคือรายการแกน และรายการรหัสอุปกรณ์ (ไม่บังคับ) ที่ระบุลําดับอุปกรณ์

หากรายการแกนว่างเปล่า เมชจะมีแกนที่ไม่มีชื่อโดยนัยขนาด 1 ในกรณีนี้ หากไม่ได้ระบุรายการรหัสอุปกรณ์ รายการรหัสอุปกรณ์โดยนัยจะเป็น [0] หากระบุรายการรหัสอุปกรณ์ รายการดังกล่าวต้องมีจำนวนเต็มรายการเดียวที่มีค่าไม่ติดลบ เราเรียกกรณีนี้ว่า "การชาร์ดดิ้งสูงสุด"

สําหรับกรณีที่ไม่ได้มีการแยกกลุ่มสูงสุดทั้งหมด หากระบุรายการรหัสอุปกรณ์ ผลคูณของขนาดแกนควรตรงกับจํานวนอุปกรณ์ หากไม่ได้ระบุรายการรหัสอุปกรณ์ รายการรหัสอุปกรณ์โดยนัยจะเป็น iota(ผลิตภัณฑ์(axes)) นอกจากนี้ เรายังไม่อนุญาตให้ระบุรายการรหัสอุปกรณ์ที่เหมือนกับ iota(product(axes)); เพื่อลดความซับซ้อน ในกรณีนี้ คุณไม่ควรระบุรายการรหัสอุปกรณ์

ตัวอย่างของเมชมีดังนี้

  • เมชว่างแสดงถึงเมชตัวยึดตําแหน่งที่แทนที่ระหว่างการนำไปใช้งานได้ propagation: <[]>
  • เมชที่มีแกนที่ไม่มีชื่อและรหัสอุปกรณ์ที่ชัดเจน ซึ่งมักใช้เพื่อแสดงการแยกกลุ่มสูงสุด: <[], device_ids=[3]>
  • ตาข่ายที่มี 2 แกนและรหัสอุปกรณ์โดยนัย iota(6): <["a"=2, "b"=3]>
  • ตาข่ายที่มี 2 แกนและรหัสอุปกรณ์ที่ชัดเจนซึ่งระบุลําดับอุปกรณ์: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
แกน ::llvm::ArrayRef<MeshAxisAttr>
device_ids ::llvm::ArrayRef<int64_t>

MeshAxisAttr

แกนที่มีชื่อในตาข่าย

ไวยากรณ์:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
ชื่อ ::llvm::StringRef ชื่อ
ขนาด int64_t

OpShardingRuleAttr

ระบุวิธีแบ่งพาร์ติชันการดำเนินการ

ไวยากรณ์:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  bool   # is_custom_rule
>

กฎการแยกส่วนระบุวิธีแบ่งพาร์ติชันการดำเนินการตามพร็อพเพอร์ตี้ต่างๆ ในการดำเนินการ เช่น แอตทริบิวต์ รูปร่างของโอเปอเรนด รูปร่างของผลลัพธ์ ฯลฯ ตัวอย่างเช่น

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>
%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

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

is_custom_rule อธิบายว่านี่เป็นกฎที่ผู้ใช้กำหนดไว้สำหรับการดำเนินการ stablehlo.custom_call หรือไม่ ตัวแบ่งพาร์ติชันไม่รู้วิธีแบ่งพาร์ติชันการดำเนินการเหล่านี้ ผู้ใช้จึงต้องบอกวิธี เมื่อเป็นกฎที่กําหนดเอง ระบบจะเก็บรักษากฎไว้เสมอ/จะไม่นําออก is_custom_rule ต้องเป็น "จริง" เท่านั้นสำหรับstablehlo.custom_call

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
factor_sizes ::llvm::ArrayRef<int64_t>
operand_mappings ::llvm::ArrayRef<TensorMappingAttr>
result_mappings ::llvm::ArrayRef<TensorMappingAttr>
is_custom_rule bool

SubAxisInfoAttr

ข้อมูลเกี่ยวกับวิธีที่แกนย่อยนี้ได้มาจากแกนทั้งหมด

ไวยากรณ์:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

เมื่อแยกแกนทั้งหมดออกเป็นแกนย่อย n แกน ระบบจะเปลี่ยนรูปร่างแกนเป็น [k_1,...,k_n] และแกนย่อยที่ i จะแสดงได้ด้วยผลคูณของขนาดแกนทั้งหมดทางด้านซ้าย m=prod(k_1,...,k_(i-1)) (หรือที่เรียกว่าขนาดก่อน) และขนาด k_i ดังนั้น แอตทริบิวต์ sub-axis-info จะเก็บตัวเลข 2 รายการดังกล่าวไว้และจะแสดงด้วยสัญลักษณ์ (m)k สำหรับขนาดก่อน m และขนาด k

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
pre_size int64_t
ขนาด int64_t

TensorMappingAttr

การแมปปัจจัยสําหรับมิติข้อมูลแต่ละมิติของเทนเซอร์

ไวยากรณ์:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
dim_mappings ::llvm::ArrayRef<DimMappingAttr>

TensorShardingAttr

การแยกกลุ่ม Tensor

ไวยากรณ์:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

การแยกกลุ่มเทนเซอร์จะเชื่อมโยงกับเมชที่เฉพาะเจาะจง และอ้างอิงชื่อแกนจากเมชนั้นได้เท่านั้น การแยกกลุ่มมิติข้อมูลจะบอกเราสำหรับมิติข้อมูลแต่ละรายการของเทนเซอร์ว่ามีการแยกกลุ่มตามแกน (หรือแกนย่อย) ใดจากแกนหลักไปแกนรอง แกนอื่นๆ ทั้งหมดที่ไม่ได้ชาร์ดมิติข้อมูลมีการจำลองแบบโดยนัยหรือโดยชัดแจ้ง (หากปรากฏในรายการแกนที่จำลอง)

คุณสามารถระบุเมชที่การแยกส่วนนี้เชื่อมโยงอยู่ได้ด้วยชื่อสัญลักษณ์ ซึ่งอ้างอิงสัญลักษณ์ MeshOp ที่เกี่ยวข้อง หรือ MeshAttr ที่ฝังอยู่

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
mesh_or_ref ::mlir::Attribute แอตทริบิวต์เมชหรือแอตทริบิวต์ข้อมูลอ้างอิงสัญลักษณ์เมชแบบแบน
dim_shardings ::llvm::ArrayRef<DimensionShardingAttr>
replicated_axes ::llvm::ArrayRef<AxisRefAttr> รายการการอ้างอิงแกน

TensorShardingPerValueAttr

การแยกกลุ่ม Tensor ตามออปเรอเรนด์/ผลลัพธ์ของการดำเนินการ

ไวยากรณ์:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

พารามิเตอร์ ได้แก่

พารามิเตอร์ ประเภท C++ คำอธิบาย
การจัดกลุ่ม ::llvm::ArrayRef<TensorShardingAttr>

Enum

PropagationDirection

enum ทิศทางการขยาย

กรณี

สัญลักษณ์ ค่า สตริง
ไม่มี 0 ไม่มี
FORWARD 1 FORWARD
BACKWARD 2 BACKWARD
ทั้งสองฝ่าย 3 ทั้งสองฝ่าย