'sdy' ภาษาถิ่น

ภาษาถิ่น Shardy (SDY)

ภาษา Shardy (SDY) กำหนดการแสดงผลการแบ่งพาร์ติชันเทนเซอร์ตามแกน และคอมโพเนนต์ API เพิ่มเติมเพื่อแนบการแบ่งพาร์ติชันกับเทนเซอร์

บันทึกเวอร์ชัน 0.0.1: เพิ่มแกนที่ไม่ได้ลดลงใน TensorShardingAttr

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

sdy.all_gather (sdy::AllGatherOp)

ดำเนินการสื่อสารแบบ All-Gather ตามแกน

ไวยากรณ์:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

รวบรวมกลุ่มของเทนเซอร์ตามแกนที่ระบุใน gathering_axes

gathering_axes คือรายการของรายการแกน รายการด้านนอกอยู่เหนือ มิติข้อมูลของ Tensor แต่ละลิสต์ด้านในจะระบุแกนที่จะใช้ดำเนินการรวบรวมแยกต่างหากในมิติข้อมูลที่เกี่ยวข้อง โดยจะใช้กับ Sharding ของตัวถูกดำเนินการ (tensor) เพื่อรับ Sharding ของผลลัพธ์ (out_sharding)

โปรดทราบว่าระบบไม่ได้ใช้ out_sharding ในการกำหนดการแบ่งส่วนของผลลัพธ์ แต่การแบ่งกลุ่มผลลัพธ์จะกำหนดโดยการแบ่งกลุ่มของ ตัวถูกดำเนินการและ gathering_axes และ out_sharding ต้องตรงกับการแบ่งกลุ่มที่อนุมานนี้

ตัวอย่าง

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• องค์ประกอบใน gathering_axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• การใช้ gathering_axes กับการแบ่งข้อมูลตัวถูกดำเนินการจะทำให้ได้ out_sharding

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.all_reduce (sdy::AllReduceOp)

Perform an all-reduce comunication along axes

ไวยากรณ์:

operation ::= `sdy.all_reduce` ($reduction_op^)? $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

ลดขนาดก้อนของเทนเซอร์ตามแกนที่ระบุใน reduction_axes ลำดับของ reduction_axes ไม่สำคัญต่อผลลัพธ์ แต่สามารถ ส่งผลต่อลำดับของกลุ่มรีพลิเคชันที่เกี่ยวข้องได้

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• reduction_axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• ต้องจัดเรียง reduction_axes ตามตาข่าย
• การแบ่งพาร์ติชันตัวถูกดำเนินการและ out_sharding ต้องมีการแบ่งพาร์ติชันมิติข้อมูลที่เทียบเท่ากัน
• reduction_axes ต้องไม่ซ้อนทับกับการแบ่งมิติข้อมูลตัวถูกดำเนินการและ แกนที่จำลอง (ซ้อนทับกับแกนที่ไม่ได้ลดได้)
• reduction_axes ต้องไม่ทับซ้อนกับแกนที่ไม่ได้ลดของ out_sharding กล่าวอีกนัยหนึ่งคือ out_sharding ต้องทำซ้ำตาม reduction_axes (โดยนัยหรือโดยชัดแจ้ง)

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.all_slice (sdy::AllSliceOp)

ดำเนินการ Slice แบบไดนามิกตามแกน

ไวยากรณ์:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

แบ่งกลุ่มของเทนเซอร์ตามแกนที่ระบุใน slicing_axes มีความสัมพันธ์แบบคู่ภาคีเชิงพีชคณิตระหว่าง sdy.all_slice กับ sdy.all_gather

slicing_axes คือรายการของรายการแกน รายการด้านนอกอยู่เหนือ มิติข้อมูลของ Tensor แต่ละรายการด้านในจะระบุแกนที่จะใช้ในการ ตัดชิ้นมิติข้อมูลที่เกี่ยวข้อง โดยจะใช้กับการแบ่งกลุ่มของตัวถูกดำเนินการ (tensor) เพื่อให้ได้การแบ่งกลุ่มของผลลัพธ์ (out_sharding)

โปรดทราบว่าระบบไม่ได้ใช้ out_sharding ในการกำหนดการแบ่งส่วนของผลลัพธ์ แต่การแบ่งกลุ่มผลลัพธ์จะกำหนดโดยการแบ่งกลุ่มของ ตัวถูกดำเนินการและ slicing_axes และ out_sharding ต้องตรงกับการแบ่งกลุ่มที่อนุมานนี้

ตัวอย่าง

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• องค์ประกอบใน slicing_axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• การใช้ slicing_axes กับการแบ่งข้อมูลตัวถูกดำเนินการจะทำให้ได้ out_sharding

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.all_to_all (sdy::AllToAllOp)

ดำเนินการสื่อสารแบบทั้งหมดต่อทั้งหมดตามแกน

ไวยากรณ์:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

สำหรับแต่ละทูเพิล (axes, src_dim, tgt_dim) ในรายการพารามิเตอร์ การดำเนินการนี้จะแบ่งกลุ่มของ Tensor ตามมิติข้อมูล tgt_dim และแกนที่ระบุใน axes กระจายกลุ่มเหล่านั้นตามแกน และต่อกันตามมิติข้อมูล src_dim

การดำเนินการนี้เป็นการรวมกันของ all-gather ตาม src_dim และ axes ตามด้วย all-slice ตาม tgt_dim และ axes กล่าวคือ ต่อท้ายมิติข้อมูลการแบ่งส่วนแกน src_dim ในเทนเซอร์อินพุต จะต่อท้ายมิติข้อมูลการแบ่งส่วนแกน tgt_dim ในเทนเซอร์เอาต์พุต

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

โปรดทราบว่าระบบไม่ได้ใช้ out_sharding ในการกำหนดการแบ่งส่วนของผลลัพธ์ แต่การแบ่งกลุ่มของผลลัพธ์จะกำหนดโดยการแบ่งกลุ่มของ ตัวถูกดำเนินการ src_dim, tgt_dim และ axes และ out_sharding ต้องตรงกับ การแบ่งกลุ่มที่สรุปนี้

ตัวอย่าง

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• ต้องระบุรายการพารามิเตอร์
• สำหรับพารามิเตอร์แต่ละรายการใน params ให้ทำดังนี้
  • องค์ประกอบใน axes ต้องเป็นไปตามข้อจำกัดของ AxisRefAttr
  • src_dim และ tgt_dim ต้องเป็นมิติข้อมูลที่ถูกต้อง (ไม่ใช่ค่าลบและน้อยกว่า อันดับของเทนเซอร์)
  • src_dim หรือ tgt_dim ต้องไม่ซ้ำกันในพารามิเตอร์ทั้งหมด
  • src_dim ต้องจัดเรียงจากน้อยไปมากในพารามิเตอร์ทั้งหมด
• การย้าย axes จาก src_dim ไปยัง tgt_dim ในการแบ่งส่วนตัวถูกดำเนินการจะทำให้ได้ out_sharding

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.collective_permute (sdy::CollectivePermuteOp)

ดำเนินการสื่อสารแบบ collective-permute เพื่อแทนที่แกน

ไวยากรณ์:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

ส่งกลุ่มของเทนเซอร์อินพุตจากแต่ละอุปกรณ์ไปยังอีกอุปกรณ์หนึ่งเพื่อ จัดลำดับใหม่/แทนที่แกนที่แยกส่วนเทนเซอร์

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

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

ในตัวอย่างด้านล่าง ขนาดของเทนเซอร์ที่แยกส่วนคือ tensor<1x4x2xf32> และ การสับเปลี่ยนแบบกลุ่มจะรักษาขนาดนี้ไว้

ตัวอย่าง

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

ข้อจำกัด:

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

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

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{}

แอตทริบิวต์

ผลลัพธ์:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

การดำเนินการที่ขอบของโฟลว์ข้อมูล

ไวยากรณ์:

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

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

Op สามารถมีขอบโฟลว์ข้อมูลหลายรายการที่ตั้งฉากกัน

เช่น

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

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

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

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

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

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

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

ลักษณะ: SameOperandsAndResultType

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

แอตทริบิวต์

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

ผลลัพธ์:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

การดำเนินการที่ขอบของโฟลว์ข้อมูลอินพุต/เอาต์พุตของฟังก์ชัน

ไวยากรณ์:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

การดำเนินการขอบของโฟลว์ข้อมูล แต่สําหรับอาร์กิวเมนต์ฟังก์ชันหรือผลลัพธ์การเรียก เมื่อตัวถูกดำเนินการเป็น BlockArgument จะเป็นสะพานจากอาร์กิวเมนต์ callOp ของผู้เรียกไปยังผู้ใช้ของอาร์กิวเมนต์ func มีขอบโฟลว์ข้อมูลฟังก์ชัน 1 รายการ สําหรับอาร์กิวเมนต์ฟังก์ชันแต่ละรายการ เมื่อตัวถูกดำเนินการเป็น OpResult จะเป็นบริดจ์ จากค่าที่ส่งคืนของ funcOp ที่เรียกไปยังผู้ใช้ผลการเรียก มีขอบโฟลว์ข้อมูลฟังก์ชัน 1 รายการสำหรับผลการเรียกแต่ละรายการ

ลักษณะ: SameOperandsAndResultType

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

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

ผลลัพธ์:

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)

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

โดยที่ body เป็นแบบโลคัลเทียบกับ manual_axes การแพร่กระจายจะเกิดขึ้นผ่าน บอดี้บนแกนอิสระใดๆ ซึ่งไม่ได้อยู่ในรายการ manual_axes

โปรดทราบว่าเทนเซอร์ที่ไม่ได้จัดอันดับควรมีการแบ่งกลุ่มที่มีอันดับเป็น 0 กล่าวคือมีการจำลองแบบเต็ม

ข้อจำกัด:

• องค์ประกอบใน in_shardings และ out_shardings ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน TensorShardingAttr
• จำนวนอินพุต/เอาต์พุตของเทนเซอร์ส่วนกลางและเทนเซอร์ภายในของภูมิภาค Op ต้องตรงกัน
• แกนแบบกำหนดเองต้องอยู่ก่อนแกนอิสระในแต่ละการแบ่งกลุ่มมิติข้อมูล
• แกนแบบกำหนดเองจะเพิ่มระยะเว้นไม่ได้ กล่าวคือ ขนาดของมิติข้อมูลต้องหารด้วยขนาดแกนที่กำหนดเองที่เกี่ยวข้องได้
• อาร์กิวเมนต์/ผลลัพธ์ของรูปร่างส่วนกลางและรูปร่างในเครื่องของภูมิภาคการดำเนินการต้องตรงกัน

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

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

แอตทริบิวต์

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

ผลลัพธ์:

sdy.mesh (sdy::MeshOp)

ตาข่ายที่มีชื่อ

ไวยากรณ์:

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

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

ลักษณะ: HasParent<ModuleOp>, SymbolName

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

แอตทริบิวต์

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

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

ตัวอย่าง

%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, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

การทำงานของแผงกั้นการแพร่กระจาย

ไวยากรณ์:

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

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

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

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

ลักษณะ: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

ผลกระทบ: MemoryEffects::Effect{}

แอตทริบิวต์

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

ผลลัพธ์:

sdy.reduce_scatter (sdy::ReduceScatterOp)

ดำเนินการสื่อสารแบบลดการกระจายตามแกน

ไวยากรณ์:

operation ::= `sdy.reduce_scatter` ($reduction_op^)? $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

ลดขนาดก้อนของเทนเซอร์ตามแกนที่ระบุใน reduce_scatter_axes แล้วกระจายผลลัพธ์ตามแกนเดียวกัน การดำเนินการนี้ โดยพื้นฐานแล้วเป็นการรวมกันของ sdy.all_reduce ตามด้วย sdy.all_slice ใน reduce_scatter_axes เดียวกัน

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• องค์ประกอบใน reduce_scatter_axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• การใช้ reduce_scatter_axes กับการแบ่งข้อมูลตัวถูกดำเนินการจะทำให้ได้ out_sharding

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

ย้ายแกนที่จำลองแบบโดยนัยหรือโดยชัดแจ้งไปยังแกนที่ไม่ได้ลด

ไวยากรณ์:

operation ::= `sdy.replicated_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes ควรจำลองในตัวถูกดำเนินการโดยนัยหรืออย่างชัดเจน การดำเนินการนี้จะทำให้ค่าดังกล่าวไม่ลดลงในผลลัพธ์ เรามีความสัมพันธ์ดังนี้

all-reduce(replicated-to-unreduced(x, axes), axes) = x

ตัวอย่าง

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}, {}\], replicated={"c", "d"}, unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.replicated_to_unreduced {"a", "c", "f"} %1 out_sharding=<@mesh, [{"b"}, {}, {}\], replicated={"d"}, unreduced={"a", "c", "e", "f"}> : tensor<8x8x8xf32>

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• ต้องจัดเรียง axes ตามตาข่าย
• axes ไม่ว่างเปล่า
• การแบ่งข้อมูลอินพุตและเอาต์พุตต้องมีการแบ่งข้อมูลมิติข้อมูลเดียวกัน
• axes ต้องจำลองโดยนัยหรือโดยชัดแจ้งในการแบ่งส่วนตัวถูกดำเนินการ
• inUnreducedAxes + axes = outUnreducedAxes

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.reshard (sdy::ReshardOp)

เปลี่ยนชาร์ดดิ้งของ Tensor เป็นชาร์ดดิ้งอื่น

ไวยากรณ์:

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

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

ทั้ง ShardingConstraintOp และ ReshardOp จะแนบการแบ่งกลุ่มกับเทนเซอร์ อายุการใช้งานของคุกกี้มีดังนี้

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

ลักษณะ: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

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

ผลกระทบ: MemoryEffects::Effect{}

แอตทริบิวต์

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

ผลลัพธ์:

sdy.return (sdy::ReturnOp)

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

ไวยากรณ์:

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

ลักษณะ: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

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

ผลกระทบ: MemoryEffects::Effect{}

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

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

ย้ายแกนที่แยกส่วนของตัวถูกดำเนินการไปยังแกนที่ไม่ได้ลดของผลลัพธ์

ไวยากรณ์:

operation ::= `sdy.sharded_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

ควรใช้ axes เพื่อแยกส่วนตัวถูกดำเนินการ การดำเนินการนี้จะทำให้ค่าเหล่านั้น ไม่ลดในผลลัพธ์ เรามีความสัมพันธ์ดังนี้

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes) โดยที่ all-gather, sharded-to-unreduced, all-reduce จะใช้กับแกนเดียวกัน

ตัวอย่าง

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\], unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.sharded_to_unreduced [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\], unreduced={"b", "c", "d", "e"}> : tensor<8x8x8xf32>

ข้อจำกัด:

• ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน Sdy_CollectiveOpInterface
• องค์ประกอบใน axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• การใช้ axes กับการแบ่งข้อมูลตัวถูกดำเนินการจะทำให้ได้ out_sharding

ลักษณะ: SameOperandsAndResultType

อินเทอร์เฟซ: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

แอตทริบิวต์

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

ผลลัพธ์:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

จำกัด Tensor ให้เป็นชาร์ดดิ้งที่ระบุ

ไวยากรณ์:

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

แนบการแบ่งพาร์ติชันกับเทนเซอร์กลาง (เช่น ผลลัพธ์ของ MatMul) เพื่อระบุว่าควรแบ่งพาร์ติชันเทนเซอร์นั้นหรือชุดย่อยของการใช้งานอย่างไร

หากการแบ่งพาร์ติชันมีมิติข้อมูลที่เปิดอยู่และแกนที่ไม่มีข้อจำกัด แสดงว่า เทนเซอร์สามารถแบ่งพาร์ติชันเพิ่มเติมตามมิติข้อมูลที่เปิดอยู่ได้

โดยการดำเนินการนี้จะทำได้ 2 วิธี ดังนี้

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

ลักษณะ: SameOperandsAndResultType

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

แอตทริบิวต์

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

ผลลัพธ์:

sdy.sharding_group (sdy::ShardingGroupOp)

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

ไวยากรณ์:

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

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

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

แอตทริบิวต์

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

Attributes

AllToAllParamAttr

พารามิเตอร์แบบทั้งหมดต่อทั้งหมด

ไวยากรณ์:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

ทูเพิลที่มีแกนและมิติข้อมูลต้นทาง/เป้าหมายเพื่อดำเนินการ แบบทั้งหมดต่อทั้งหมด

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

AllToAllParamListAttr

รายการพารามิเตอร์แบบทั้งหมดต่อทั้งหมด

ไวยากรณ์:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

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

AxisRefAttr

การอ้างอิงถึงแกนแบบเต็มหรือแกนย่อยที่แยก

ไวยากรณ์:

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

ข้อจำกัด:

• name ต้องอยู่ใน MeshAttr ที่ผูกไว้
• หากมี sub_axis_info จะต้องเป็นไปตามข้อจำกัดของ SubAxisInfoAttr

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

AxisRefListAttr

รายการอ้างอิงแกน

ไวยากรณ์:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

ข้อจำกัด:

• องค์ประกอบใน value ต้องเป็นไปตามข้อจำกัดของ AxisRefAttr
• ไม่มีการอ้างอิงแกนหรือแกนย่อยที่ซ้ำกันซึ่งทับซ้อนกัน
• ไม่มีการอ้างอิงแกน 2 รายการที่อยู่ติดกันซึ่งเป็นแกนย่อยที่ต่อเนื่องของแกนแบบเต็มเดียวกัน กล่าวคือ สามารถผสานรวมเป็นแกนย่อยหรือแกนแบบเต็มได้

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

AxisToPropagationDetailsAttr

รายละเอียดโฟลว์ขอบการเผยแพร่สำหรับแกนและแหล่งที่มาที่เฉพาะเจาะจง

ไวยากรณ์:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

จับคู่การอ้างอิงค่าแหล่งที่มากับรายการการอ้างอิงค่าเป้าหมายตามแกนที่เฉพาะเจาะจง

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

DimMappingAttr

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

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

ข้อจำกัด:

• มีดัชนีปัจจัยอย่างน้อย 1 รายการ
• ดัชนีปัจจัยต้องอยู่ในช่วง [0, $factor_sizes)
• หากมีหลายปัจจัย จะต้องไม่มีปัจจัยใดมีขนาดเป็น 1
• ไม่มีดัชนีปัจจัยที่ซ้ำกัน

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

DimensionShardingAttr

การแบ่งมิติข้อมูล

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

ข้อจำกัด:

• องค์ประกอบใน axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• หากการแยกมิติข้อมูลมีลำดับความสำคัญ ให้ทำดังนี้
  • ลำดับความสำคัญมากกว่าหรือเท่ากับ 0
  • มิติข้อมูลมีแกนอย่างน้อย 1 แกนหากปิดอยู่

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

EdgeValueRefAttr

การอ้างอิงถึงดัชนีที่เฉพาะเจาะจงของขอบค่าประเภท type

ไวยากรณ์:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

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

ListOfAxisRefListsAttr

รายการอ้างอิงของแกน

ไวยากรณ์:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

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

ManualAxesAttr

รายการแกนที่ ManualComputationOp เป็นแบบกำหนดเอง

ไวยากรณ์:

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

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

MeshAttr

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

ไวยากรณ์:

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

Mesh คือรายการแกนและรายการรหัสอุปกรณ์ที่ไม่บังคับซึ่งระบุลำดับอุปกรณ์

หากรายการแกนว่างเปล่า

• หากไม่ได้ระบุ device_ids จะถือว่าเป็นตาข่ายว่าง
• หากระบุ device_ids จะต้องเป็นจำนวนเต็มที่ไม่ใช่ค่าลบ รายการเดียว ซึ่งเราเรียกว่าตาข่ายการแยกส่วนสูงสุด

หากระบุรายการแกน

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

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

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

ข้อจำกัด:

• องค์ประกอบใน device_ids ต้องไม่เป็นค่าลบ
• หาก axes ว่างเปล่า ขนาดของ device_ids จะเป็น 0 (ตาข่ายว่างเปล่า) หรือ 1 (ตาข่ายการแยกข้อมูลสูงสุด) ก็ได้
• หาก axes ไม่ว่างเปล่า ให้ทำดังนี้
  • องค์ประกอบใน axes ต้องไม่มีชื่อซ้ำกัน
  • หากระบุ device_ids แล้ว device_ids เดิมจะไม่ใช่ iota(product(axis_sizes)) และ device_ids ที่จัดเรียงแล้วจะเป็น iota(product(axis_sizes))

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

MeshAxisAttr

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

ไวยากรณ์:

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

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

OpShardingRuleAttr

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

ไวยากรณ์:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  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 แม้ว่าจะไม่สามารถแยกส่วนได้ ซึ่งส่วนใหญ่เป็นไปเพื่อให้สมบูรณ์เนื่องจาก Op หลายรายการ เช่น Op แบบ Pointwise มีขนาด มิติข้อมูล 1 ที่สอดคล้องกันในตัวถูกดำเนินการและผลลัพธ์

ประเภทปัจจัย:

• reduction_factors มีดัชนีของแฟกเตอร์ที่ต้องลด เช่น มิติข้อมูลการทำสัญญาในการดำเนินการแบบจุด ปัจจัยเหล่านี้อาจอยู่ในตัวถูกดำเนินการแต่ไม่อยู่ในผลลัพธ์
• need_replication_factors มีดัชนีของปัจจัยที่ต้องมีการจำลองแบบเต็ม เช่น มิติข้อมูลที่จัดเรียงในการดำเนินการจัดเรียง
• permutation_factors มีดัชนีของปัจจัยที่ต้องใช้ collective-permute หากมีการแบ่งส่วน เช่น มิติข้อมูลการเพิ่มพื้นที่ใน การดำเนินการ pad
• ปัจจัยอื่นๆ ทั้งหมดถือเป็นปัจจัยแบบส่งผ่าน ซึ่งหมายถึงปัจจัย ที่ไม่ต้องมีการสื่อสารใดๆ หากมีการแบ่งข้อมูลในลักษณะเดียวกันในเทนเซอร์ทั้งหมด ที่แมปกับปัจจัยเหล่านั้น

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

is_custom_rule อธิบายว่านี่คือกฎที่ผู้ใช้กำหนดหรือไม่ ผู้ใช้ สามารถกำหนดกฎการแบ่งข้อมูลสำหรับคอลที่กำหนดเองหรือเขียนทับ กฎการแบ่งข้อมูลที่กำหนดไว้ล่วงหน้าสำหรับการดำเนินการมาตรฐาน ระบบจะเก็บรักษากฎที่กำหนดเองไว้เสมอ/ไม่นำออก

ข้อจำกัด:

• จำนวนการแมปตัวถูกดำเนินการ/ผลลัพธ์ต้องตรงกับจำนวนตัวถูกดำเนินการ/ผลลัพธ์ของ Op
• มีการแมปอย่างน้อย 1 รายการ (ต้องมีกฎสำหรับการดำเนินการที่ไม่มีตัวถูกดำเนินการ/ผลลัพธ์)
• แรงค์ของแต่ละ TensorMappingAttr จะตรงกับแรงค์ของ ประเภทเทนเซอร์ที่เกี่ยวข้อง
• สำหรับแต่ละกลุ่มของปัจจัย (reduction_factors, need_replication_factors, permutation_factors) ให้ทำดังนี้
  • องค์ประกอบต้องอยู่ในช่วง [0, $factor_sizes]
  • ไม่มีดัชนีปัจจัยที่ซ้ำกันภายในแต่ละกลุ่มและในกลุ่มต่างๆ

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

PropagationEdgesAttr

ข้อมูลเมตาของ Edge การเผยแพร่สำหรับขั้นตอนการเผยแพร่ทั้งหมด

ไวยากรณ์:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

รายการรายละเอียดการแพร่กระจายต่อแกนสำหรับค่า โดยจัดกลุ่มตามดัชนีขั้นตอน

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

PropagationOneStepAttr

ข้อมูลเมตาการเผยแพร่ต่อขั้นตอน

ไวยากรณ์:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

รายละเอียดการขยายสำหรับแกนทั้งหมดในขั้นตอนการขยายเดียว

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

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

ข้อจำกัด:

• pre-size มีค่าอย่างน้อย 1
• size มากกว่า 1
• pre-size ต้องหารขนาดของแกนเต็ม นั่นคือทั้ง pre-size และ size หารขนาดของแกนเต็ม และแกนย่อยต้องไม่เกินแกนเต็ม
• ขนาดของแกนย่อยไม่เท่ากับขนาดของแกนเต็มที่สอดคล้องกัน ในกรณีนี้ควรใช้แกนเต็มแทน

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

TensorMappingAttr

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

ไวยากรณ์:

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

ข้อจำกัด:

• องค์ประกอบใน dim_mappings ต้องเป็นไปตามข้อจำกัดใน DimMappingAttr
• ไม่มีดัชนีปัจจัยที่ซ้ำกันในมิติข้อมูล

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

TensorShardingAttr

การแบ่งส่วนเทนเซอร์

ไวยากรณ์:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

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

โปรดทราบว่าการไม่มีแอตทริบิวต์การแยกส่วนในเทนเซอร์จะเทียบเท่ากับการแยกส่วนเทนเซอร์แบบเปิดอย่างเต็มรูปแบบ

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

การแยกส่วนอาจมีแกนที่ไม่ได้ลด (ระบุโดย unreduced_axes) ซึ่งหมายความว่า เทนเซอร์ไม่ได้ลดตามแกนเหล่านี้ เช่น หากมิติข้อมูลการทำสัญญาของ matmul มีการแบ่งพาร์ติชันตามแกน x ทั้งใน lhs และ rhs ผลลัพธ์จะไม่มีการลดตาม x การใช้ all-reduce กับ Tensor ตามแกนที่ไม่ได้ลดจะทำให้ Tensor จำลองตามแกนเหล่านั้น อย่างไรก็ตาม Tensor ที่มีแกนที่ไม่ได้ลดไม่จำเป็นต้องลดทั้งหมดทันที แต่สามารถคงสถานะไม่ได้ลดไว้เมื่อส่งไปยังการดำเนินการเชิงเส้น เช่น stablehlo.add (ตราบใดที่ทั้ง lhs และ rhs ไม่ได้ลด) และลดทั้งหมดในภายหลัง เราถือว่าประเภทการลดคือผลรวม ส่วนการลดอื่นๆ อาจ รองรับในอนาคต

ข้อจำกัด:

• องค์ประกอบใน dim_shardings ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน DimensionShardingAttr
• องค์ประกอบใน replicated_axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• องค์ประกอบใน unreduced_axes ต้องเป็นไปตามข้อจำกัดที่ระบุไว้ใน AxisRefListAttr
• หากประเภท Tensor ที่เกี่ยวข้องไม่ใช่ ShapedType การชาร์ดดิ้งต้องมีอันดับเป็น 0 และไม่มีแกนที่จำลอง
• หากเป็น ShapedType ให้ทำดังนี้
  • เทนเซอร์ควรมีอันดับ
  • จำนวนการแบ่งมิติข้อมูลเท่ากับอันดับของเทนเซอร์
  • ระบบจะไม่แยกมิติข้อมูลที่มีขนาด 0
• ไม่มีการอ้างอิงแกนหรือแกนย่อยที่ซ้ำกันซึ่งทับซ้อนกัน ใน dim_shardings, replicated_axes และ unreduced_axes
• รายการใน replicated_axes และ unreduced_axes จะเรียงตาม mesh_or_ref (ดู AxisRefAttr::getMeshComparator)

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

TensorShardingPerValueAttr

การแยกส่วน Tensor ต่อตัวถูกดำเนินการ/ผลลัพธ์ของ Op

ไวยากรณ์:

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

รายการของ TensorShardingAttr โดยมี 1 รายการสำหรับตัวถูกดำเนินการ/ผลลัพธ์แต่ละรายการของ Op

ข้อจำกัด:

• องค์ประกอบใน shardings ต้องเป็นไปตามข้อจำกัดของ TensorShardingAttr

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

Enum

EdgeNodeType

Edge node type enum

เคส

PropagationDirection

การแจงนับทิศทางการเผยแพร่

เคส

ReductionOp

Reduction op enum

เคส