-mpmd-copy-topology-from-main

คัดลอกโทโพโลยีจากฟังก์ชันหลักไปยังฟังก์ชันที่ mpmd.call อ้างอิง

คัดลอกแอตทริบิวต์โทโพโลยีจากฟังก์ชันหลักไปยังฟังก์ชันใดก็ตาม ที่อ้างอิงโดย mpmd.call นอกจากนี้ยังตั้งค่า mpmd.call callee เป็น private visibility เพื่อไม่ให้เข้าใจผิดว่าเป็นฟังก์ชันจุดแรกเข้า

-mpmd-enforce-input-output-equisharding

บังคับใช้ข้อจํากัดการแบ่งข้อมูลเท่ากันสําหรับฟังก์ชัน MPMD

บังคับใช้ข้อจำกัดการแบ่งพาร์ติชันที่เท่ากันของอินพุตและเอาต์พุตสำหรับฟังก์ชัน MPMD โดย การเปิดตัว TransferOps เมื่อจำเป็น

ตัวเลือก

-constraints : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.

-mpmd-generate-sdy-meshes-from-topology

สร้างเมชแบบแยกส่วนตามโทโพโลยี MPMD

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

-mpmd-infer-mesh-assign-mesh-func-leaves

_กำหนดตาข่ายให้กับการคำนวณ เอาต์พุตฟังก์ชัน และอินพุตฟังก์ชันที่ไม่ได้ใช้แต่ละรายการโดยใช้การวิเคราะห์ use_set และ srcset

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

นอกจากนี้ เรายังถือว่าค่ากลางบางค่าเป็นค่าใบไม้เพื่อประโยชน์ในการวิเคราะห์ด้วย กล่าวคือ ตัวถูกดำเนินการของ mpmd.reduce และ mpmd.broadcast จะถือเป็นใบ และระบบจะสร้างคู่การกำหนด-ยกเลิกการกำหนดในตัวถูกดำเนินการเหล่านั้น

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

การส่งผ่านนี้จะล้มเหลวและแสดงข้อผิดพลาดหากไม่ได้ป้อนข้อมูล use- และ src-sets อย่างถูกต้อง สำหรับ Leaf Op

เงื่อนไขเบื้องต้น: ทุกการดำเนินการมี src-set ที่ไม่ว่างเปล่า หรือเราอนุมานการโอน

ตัวเลือก

-infer-transfers : Whether to create transfers when needed, instead of erroring.
-error-limit     : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-infer-mesh-assign-using-input-output-constraints

กำหนดตาข่ายให้กับอินพุตและเอาต์พุตตามข้อจำกัดการกำหนดอินพุต-เอาต์พุต

การส่งผ่านนี้ใช้ข้อจำกัดการกำหนดค่าที่เท่ากันของอินพุตและเอาต์พุตเพื่อกำหนดทั้งอินพุตและเอาต์พุตให้กับ Mesh เดียวกัน

โปรดทราบว่าการดำเนินการนี้จะรับประกันว่าอินพุตอยู่ใน Mesh เดียวกัน ไม่ว่าจะมีการโอนไปยัง Mesh อื่นๆ ในภายหลังหรือไม่ก็ตาม แต่หมายความว่า เราไม่ควรเรียกใช้ populate-src-set หลังจากผ่านขั้นตอนนี้

ต้องมี

  • สำหรับอินพุต i ของฟังก์ชันจุดแรกเข้าซึ่งอาจเป็นส่วนหนึ่งของข้อจํากัดการกำหนดค่าที่เท่ากัน: i มี MeshTensorType หรือชุดการใช้งานที่กำหนดไว้อย่างชัดเจน
  • สำหรับเอาต์พุต o ของฟังก์ชันจุดแรกเข้าที่อาจเป็นส่วนหนึ่งของข้อจำกัดการกำหนดค่าเท่ากัน o มีประเภท MeshTensorType หรือมีชุด src และ use ที่กำหนดไว้อย่างดี

โดยที่ชุดการใช้งานที่กำหนดไว้อย่างดีของค่าจะรวมตาข่ายทั้งหมดที่ค่า (แบบทรานซิทีฟ) กำหนดให้ผ่านการดำเนินการ mpmd.assign และไม่มีตาข่ายอื่นๆ src-set ที่กำหนดไว้อย่างดีจะมีตาข่ายทั้งหมดที่อนุญาตให้เทนเซอร์อยู่ได้และไม่มีตาข่ายอื่นๆ

แม้ว่าฟังก์ชันนี้จะทำงานเฉพาะในฟังก์ชันจุดแรกเข้า แต่เราก็ทำให้เป็นโมดูล Op pass เนื่องจากต้องมีการส่งผ่านทั้งหมดในฟังก์ชันที่มีอยู่ให้เสร็จสมบูรณ์ ก่อนจึงจะทำงานได้ เช่น หากเราทำให้เป็น EntryPointFunctionPass ตัวจัดการบัตรอาจเรียกใช้บัตรนี้ก่อนที่การตรวจสอบจะเสร็จสมบูรณ์ในฟังก์ชันที่ไม่ใช่จุดแรกเข้า

ตัวเลือก

-verbose-logging : Whether to enable verbose logging
-constraints     : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.

-mpmd-infer-mesh-convert-reduce-ops

แปลงการดำเนินการ reduce ที่มีคำอธิบายประกอบเป็นการดำเนินการ mpmd.reduce และทำให้การดำเนินการ reduce เป็นแบบแบน

แปลงการดำเนินการ Reduce ที่มีคำอธิบายประกอบเป็นการดำเนินการ mpmd.reduce และทำให้การดำเนินการ Reduce เหล่านี้แบนราบ

ในสัญลักษณ์

x = add(w0, w1) {mpmd.reduce = #mpmd.reduce} y = add(x, w2) {mpmd.reduce = #mpmd.reduce} ~~> r = mpmd.reduce(w0,w1,w2)

ตัวเลือก

-infer-cross-mesh-reductions : Whether to infer cross-mesh reductions. Will be enabled by default once stable.

-mpmd-infer-mesh-finalize

ใช้การล้างข้อมูลขั้นสุดท้ายหลังจากอนุมานรูปแบบตาข่าย

ตัวเลือก

-infer-transfers : Whether to create transfers when needed, instead of erroring.

-mpmd-infer-mesh-populate-src-set

_Initializes the src_set for UnassignOps and func args and propagates the srcset.

การส่งผ่านนี้จะเริ่มต้น src_set และเผยแพร่ src_set โดยการป้อนข้อมูลกราฟ ด้วยข้อมูล src_set

เงื่อนไขเบื้องต้น: หากต้องการให้ฟังก์ชันอาร์กิวเมนต์มี src_sets คุณต้องป้อนข้อมูล use_set

การเริ่มต้น: src_set ของ UnassignOp จะตั้งค่าเป็น Mesh ที่กำหนดให้ ระบบจะตั้งค่า src_set ของอาร์กิวเมนต์ฟังก์ชันเป็น use_set

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

-mpmd-infer-mesh-populate-use-set

_Initializes the use_set for AssignOps and propagates the useset.

การส่งผ่านนี้จะเริ่มต้น use_set และเผยแพร่ย้อนกลับ โดยจะป้อนข้อมูล กราฟด้วยข้อมูล use_set

การเริ่มต้น: ตั้งค่า use_set ของ AssignOp เป็น Mesh ที่กำหนดให้

การแพร่กระจาย: use_sets จะแพร่กระจายย้อนกลับจากผู้ใช้ไปยังการดำเนินการเอง โดยใช้การรวมผู้ใช้ use_set ของการดำเนินการคือการรวม use_set ของผู้ใช้ตาม คำจำกัดความ เนื่องจาก use_set คือชุดของการใช้งานแบบทรานซิทีฟ

-mpmd-infer-mesh-rewrite-using-analysis

_เขียนการดำเนินการใหม่ตามชุดการใช้งาน

การส่งผ่านนี้จะกำหนดการดำเนินการแบบไม่มีตาข่ายโดยการห่อไว้ใน Fragment โดยใช้การวิเคราะห์ use_set และ src_set

นอกจากนี้ ยังนำแอตทริบิวต์ use_set และ src_set ออกด้วยเพื่อเป็นการล้างข้อมูล เนื่องจาก ไม่จำเป็นต้องวิเคราะห์อีกต่อไปหลังจากนี้

เงื่อนไขเบื้องต้น: ทุกการดำเนินการมี use_set ซึ่งหมายความว่าการวิเคราะห์เสร็จสมบูรณ์แล้ว เงื่อนไขเบื้องต้น: อาร์กิวเมนต์ทุกรายการของฟังก์ชันที่ไม่ใช่จุดแรกเข้าจะใช้โดย Op ที่ไม่ใช่ตัวสิ้นสุดอย่างน้อย 1 รายการ

TODO: jupvfranco - consider renaming this pass given that it doesn't depend on the analysis so much anymore.

ตัวเลือก

-max-clones : How many copies of a meshless operation we allow. Setting it to 1 means we never clone the op.

-mpmd-infer-mesh-validate-no-additional-transfers-needed

ตรวจสอบว่าไม่จำเป็นต้องโอนเพิ่มเติมสำหรับการกำหนดค่า Mesh

การทดสอบนี้จะตรวจสอบว่าการกำหนดค่า Mesh เป็นไปได้สำหรับ Ops ทั้งหมดที่ไม่มี Mesh โดยไม่ต้องมีการโอนเพิ่มเติม

สำหรับการดำเนินการแบบไม่มี Mesh ซึ่งไม่ใช่การดำเนินการฟังก์ชัน ระบบจะแสดงข้อผิดพลาดในกรณีต่อไปนี้

  1. use_set ไม่อยู่ใน src_set สำหรับการดำเนินการที่กำหนด กล่าวคือ ต้องมีการโอน

สำหรับฟังก์ชันการทำงานของฟังก์ชัน ก็เพียงพอที่จะตรวจสอบเงื่อนไขข้างต้นสำหรับอาร์กิวเมนต์ของฟังก์ชัน เนื่องจากฟังก์ชันจะแสดงผลการดำเนินการแบบไม่มีตาข่ายหรืออาร์กิวเมนต์ของบล็อก

ตัวเลือก

-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-infer-mesh-validate-src-set-not-empty

_ตรวจสอบว่าทุกการดำเนินการที่ไม่มี Mesh มี srcset ที่ไม่ว่างเปล่า

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

สำหรับ Func Ops การตรวจสอบเงื่อนไขข้างต้นสำหรับ Func Args ก็เพียงพอแล้ว เนื่องจาก Func จะแสดงผลเป็น Meshless Ops หรือ Block Args

ซึ่งต้องเป็นพาสในระดับโมดูล เนื่องจากในกรณีที่เกิดข้อผิดพลาด ในผู้ถูกเรียก เราต้องการพิมพ์ผู้เรียก

เงื่อนไขเบื้องต้น: ควรแปลงการลดแบบครอสเมชเพื่อลดการดำเนินการก่อนที่จะเรียกใช้พาสนี้

ตัวเลือก

-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-inline-nested-user-exposed-ops

_แทรกการดำเนินการ mpmd ที่ผู้ใช้มองเห็นซึ่งซ้อนอยู่ในการคำนวณที่มีชื่อ

แทรก named_computation, named_tensor, broadcast และ reduce op ที่ซ้อนอยู่ใน named_computation โดยตรวจสอบว่าการกำหนดค่า Mesh (เมื่อกำหนด) ตรงกับการกำหนดค่า Mesh ของระดับบนสุด

ตัวเลือก

-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.

-mpmd-insert-nameless-clone-of-negligible-ops

การโคลนการดำเนินการที่ไม่มีนัยสำคัญนอกการคำนวณที่มีชื่อ

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

การส่งผ่านนี้จะไม่เปลี่ยนการคำนวณที่ระบุชื่อเลย

-mpmd-introduce-transfers

สร้างการโอนข้อมูลตามการกำหนดตาข่ายผู้ใช้

สร้างพาสที่แนะนำการดำเนินการโอนตามการกำหนดค่า Mesh ของผู้ใช้ ซึ่งรวมถึงเนื้อหาต่อไปนี้

  1. ส่ง UnassignOp ไปยังการเรียก mpmd หากผลลัพธ์ของ UnassignOp มีการกำหนดในผู้รับสายในภายหลัง
  2. แทนที่ AssignOp ของ UnassignOp ด้วย TransferOp
  3. กำหนดการเพิ่มลงในเมชที่ใช้ และเริ่มการโอนหากมีการเพิ่มแบบไม่มีเมชระหว่างชิ้นส่วน

-mpmd-map-input-output-to-mesh

กำหนดตาข่ายให้กับอินพุตและเอาต์พุตของฟังก์ชัน

สร้างพาสที่แมปอินพุต/เอาต์พุตของฟังก์ชันกับเมชที่กำหนด การกำหนดเมชที่ผู้ใช้กำหนด

สำหรับการอาร์กิวเมนต์อินพุต การส่งผ่านนี้จะทำดังนี้

  1. ส่งเมตริกซ์อินพุตที่ควรวางไว้ในเมชไปยังเมช เทนเซอร์
  2. อัปเดตลายเซ็นฟังก์ชัน
  3. เพิ่ม mpmd.unassign ก่อนที่จะใช้เทนเซอร์

สำหรับอาร์กิวเมนต์เอาต์พุต พาสนี้จะเพิ่ม mpmd.assign ก่อนที่จะส่งคืนเทนเซอร์ และอัปเดตลายเซ็นฟังก์ชัน

ต้องมี: ดัชนีอินพุต/เอาต์พุตแต่ละรายการถูกต้องและแต่ละเมชที่แมป เป็นเมชที่ถูกต้องในโทโพโลยี

ตัวเลือก

-input-assignment  : Mapping between function input indices and assigned mesh names.E.g., '0@m0,1@m1' defines that input with index 0 will be assigned to mesh m0 and input with index 1 will be assigned to mesh m1.
-output-assignment : Mapping between function output indices and assigned mesh names.E.g., '0@m0,1@m1' defines that output with index 0 will be assigned to mesh m0 and output with index 1 will be assigned to mesh m1.

-mpmd-map-named-ops-to-mpmd-ops

กำหนดตาข่ายให้กับการดำเนินการที่ผู้ใช้กำหนด

สร้างพาสโดยกำหนด mpmd.named_tensor ให้กับ Assign(Unassign(%v)) (ขึ้นอยู่กับว่ามีรายการใน assignment หรือไม่) และกำหนด map แต่ละ named_computation ให้กับเมช โดยใช้การแมปที่ผู้ใช้กำหนดระหว่าง named_computation กับชื่อเมช ซึ่งหมายถึงการแทนที่ named_computation แต่ละรายการด้วย Fragment และการสร้าง AssignOps สำหรับตัวถูกดำเนินการ และ UnassignOps สำหรับผลลัพธ์ของ Fragment เหล่านี้ ตอนนี้เราได้เปลี่ยนรูปแบบ Assign(Unassign(%v)) เป็น Transfer(%v) แล้ว ไม่มีการดำเนินการ named_computation/named_tensor หลังจากขั้นตอนนี้

ต้องมี: named_computations และ named_tensors ทั้งหมดอยู่ที่ระดับบนสุด ของฟังก์ชัน

ตัวเลือก

-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.

-mpmd-simplify-named-computations

ลดความซับซ้อนของอินพุตและเอาต์พุตของ Op การคำนวณที่มีชื่อ

ลดความซับซ้อนของการคำนวณที่มีชื่อแต่ละรายการอย่างอิสระ โดยเฉพาะอย่างยิ่ง

  • ขจัดผลลัพธ์ที่ซ้ำกันและค่าที่ส่งคืนที่เกี่ยวข้อง
  • ขจัดข้อมูลที่ซ้ำกันของตัวถูกดำเนินการและอาร์กิวเมนต์บล็อกที่เกี่ยวข้อง
  • นำผลลัพธ์ที่มีตัวถูกดำเนินการส่งคืนที่เกี่ยวข้องเป็นบล็อก อาร์กิวเมนต์ของตัวดำเนินการออก
  • นำตัวถูกดำเนินการที่มีอาร์กิวเมนต์บล็อกที่สอดคล้องกันซึ่งไม่มีการใช้งานอีกต่อไป (หรือ ไม่มีการใช้งานตั้งแต่แรก) ออก และ
  • นำผลลัพธ์ที่ไม่ได้ใช้ออก
  • แทนที่รูปแบบ arg -> stablehlo.optimization_barrier -> return ในการคำนวณที่มีชื่อด้วยรูปแบบ arg -> return ซึ่งช่วยให้ ลดความซับซ้อนได้มากขึ้น

-mpmd-validate-named-ops-in-mpmd-func

ตรวจสอบว่าการดำเนินการที่มีชื่อซ้อนอยู่ในฟังก์ชัน mpmd เท่านั้น

ตรวจสอบว่า NamedComputationOp และ NamedTensorOp ซ้อนกันอยู่ในฟังก์ชัน mpmd เท่านั้น ซึ่งก็คือฟังก์ชันที่มีแอตทริบิวต์โทโพโลยี