จาก HLO ไปยัง Thunk

เอกสารนี้อธิบายเส้นทางของโมดูล High Level Optimizer (HLO) ของ XLA ตั้งแต่สถานะเริ่มต้นไปจนถึงไฟล์ที่เรียกใช้งานได้ขั้นสุดท้าย บางครั้งเราจะละคำว่า "โมดูล" และเรียกเพียงว่า "HLO"

แผนภาพ HLO เป็น Thunk

HLO ก่อนการเพิ่มประสิทธิภาพ

เราเริ่มต้นด้วยโมดูล HLO ก่อนการเพิ่มประสิทธิภาพ HLO ก่อนการเพิ่มประสิทธิภาพไม่มีการดำเนินการ (ops) ที่ถือว่าเป็นภายใน XLA เช่น fusion หรือ bitcast โดยที่การดำเนินการจะไม่มีเลย์เอาต์ในขั้นตอนนี้ หรือหากมี ระบบจะ ไม่สนใจ โดยปกติแล้ว HLO ก่อนการเพิ่มประสิทธิภาพจะสร้างขึ้นโดยเฟรมเวิร์กระดับสูงกว่า เช่น TensorFlow และ JAX เมื่อใช้แฟล็ก XLA -xla_dump_to ระบบจะส่งออก HLO ก่อนการเพิ่มประสิทธิภาพไปยังไฟล์ที่มีคำต่อท้ายชื่อไฟล์เป็น “before_optimizations.txt”

เพิ่มประสิทธิภาพโมดูล HLO

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

ซึ่งรวมถึงการส่งผ่าน เช่น Shardy Partitioner หรือการส่งผ่านสำหรับการแยกส่วน SPMD

การเพิ่มประสิทธิภาพผ่าน

ซึ่งอาจรวมถึงการผ่านกฎหมายและการลดความซับซ้อน

การเพิ่มประสิทธิภาพแบบรวม

คล้ายกับการเพิ่มประสิทธิภาพ แต่เน้นที่การดำเนินการร่วมกัน

บัตรผ่านการมอบหมายเลย์เอาต์

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

ตัวอย่างรูปร่างที่มีเลย์เอาต์

f32[10,20,30]{2,0,1}

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

เป้าหมายของการกำหนดเลย์เอาต์คือการลดจำนวนการสับเปลี่ยนตำแหน่งจริงที่จำเป็นโดยใช้กลยุทธ์แบบละโมบ โดยเริ่มจากข้อจำกัดของเลย์เอาต์บางอย่าง (เช่น ไลบรารี cuDNN/cuBLAS คาดหวังมิติข้อมูลที่ต่อเนื่องกัน) และเผยแพร่เลย์เอาต์ "ลง" แล้ว "ขึ้น" กราฟ HLO เมื่อสิ้นสุดการส่งต่อเลย์เอาต์ คำสั่งบางอย่างอาจมีเลย์เอาต์ที่ขัดแย้งกัน โดยคำสั่งหนึ่งส่งต่อจากตัวถูกดำเนินการ และอีกคำสั่งหนึ่งส่งต่อจากผู้ใช้ หากต้องการแก้ไขความขัดแย้งนี้ ระบบจะแทรกคำสั่ง copy HLO ที่เปลี่ยนเลย์เอาต์จากเลย์เอาต์ของตัวถูกดำเนินการเป็นเลย์เอาต์ของคำสั่ง

การแปลงเลย์เอาต์เป็นรูปแบบมาตรฐาน

เนื่องจากรูปร่างทางกายภาพค่อนข้างยากที่จะระบุ การปรับเลย์เอาต์ให้เป็นมาตรฐานจึงพยายามเขียนรูปร่างใหม่เพื่อให้ใช้เลย์เอาต์เริ่มต้น {rank-1, rank-2, …, 0} ในตัวอย่างด้านบน รูปร่างที่ปรับให้เป็นมาตรฐานจะเป็น f32[20,10,30]{2,1,0} ระบบจะเขียนการดำเนินการคัดลอกที่เปลี่ยนเลย์เอาต์ใหม่เป็น การรวมกันของ transpose และ bitcast เนื่องจากปัจจุบันเราไม่สามารถ ทำให้การดำเนินการทั้งหมดเป็นมาตรฐานได้ จึงยังมีการดำเนินการบางอย่างที่อาจมีเลย์เอาต์ที่ไม่ใช่ค่าเริ่มต้น โดยเฉพาะอย่างยิ่ง gather และ dot ที่ขอบเขตระหว่างการดำเนินการที่ทำให้เป็นปกติและการดำเนินการที่ไม่ทำให้เป็นปกติ จะมีbitcastการดำเนินการที่แสดงถึงการเปลี่ยนตำแหน่ง นั่นคือ การเปลี่ยนตำแหน่งที่มีการกำหนดเลย์เอาต์ซึ่งทำให้ไม่มีการดำเนินการใดๆ ทางกายภาพ

การทําให้เลย์เอาต์เป็นปกติยังทําให้การสลับที่โดยนัยบางอย่างชัดเจนขึ้น ซึ่งเป็นสิ่งสําคัญเนื่องจาก Codegen สามารถจัดการการสลับที่ที่ชัดเจนด้วยเครื่องมือปล่อยสัญญาณเฉพาะได้ เช่น ในทางเทคนิคแล้ว การเปลี่ยนรูปร่างจะได้รับอนุญาตให้มีเลย์เอาต์ทางกายภาพที่แตกต่างกันระหว่างตัวถูกดำเนินการกับผลลัพธ์ (เช่น เนื่องจากอันดับที่แตกต่างกัน) พาส ReshapeDecomposerที่ทำงานเป็นส่วนหนึ่งของพาสการทำให้เลย์เอาต์เป็นปกติ จะเปลี่ยนการเปลี่ยนรูปร่างเป็นลำดับของ transpose, การเปลี่ยนรูปร่าง bitcast และ transpose

การเพิ่มประสิทธิภาพการกำหนดเลย์เอาต์หลังการโพสต์

การส่งผ่านที่สำคัญที่สุดที่นี่คือการผสาน Triton (การผสาน GEMM + การผสาน Softmax/Layernorm) หรือการเขียนใหม่เป็นการเรียกใช้ไลบรารี การปรับอัตโนมัติจะทำงานในขั้นตอนนี้ด้วย โดย XLA จะเลือกระหว่าง Emitter ต่างๆ เลือกอัลกอริทึมที่ดีที่สุดสำหรับการผสานรวมหรือการคูณเมทริกซ์ ค้นหาการแบ่งไทล์ที่ดีที่สุดสำหรับการผสานรวมที่จัดการโดย Triton Emitter เป็นต้น

บัตรฟิวชัน

โดยมี 2 พาสหลักคือ PriorityFusion และ Multi-Output ฟิวชัน

ใน PriorityFusion เราจะสร้างฟิวชันตามโมเดลต้นทุน เมื่อผสาน เราจะอนุญาตให้ทำซ้ำการดำเนินการกับผู้ใช้หลายรายหากผสานการดำเนินการกับผู้ใช้ทั้งหมดได้ นอกจากนี้ เรายังอนุญาตให้ขยายการผสาน Softmax ของ Triton ที่มีอยู่หากเป็นไปได้

Multi-Output fusion เป็นการส่งผ่านแยกต่างหากที่อนุญาตให้ผสานการดำเนินการ/การผสานที่ แชร์ตัวถูกดำเนินการ นอกจากนี้ยังผสานตัวถูกดำเนินการ/การผสานตัวถูกดำเนินการเข้ากับผู้ใช้ได้โดยไม่ต้องทำซ้ำด้วยการเพิ่มเอาต์พุตพิเศษ เพื่อให้ผู้ใช้รายอื่นๆ ของตัวดำเนินการที่จะผสานรวมสามารถเปลี่ยนเส้นทางไปยังเอาต์พุตเหล่านี้ได้ การส่งผ่านนี้ต้องระมัดระวังไม่ให้เกิด รอบในกราฟ HLO

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

การส่งผ่านหลายครั้งหลังการรวม

การส่งผ่านหลายรายการที่เกี่ยวข้องกับ Collective (เช่น การเปลี่ยนเป็นแบบไม่พร้อมกัน หรือการบังคับใช้ ลำดับที่เกี่ยวข้องของ Collective บางรายการ)

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

เมื่อสิ้นสุดการเพิ่มประสิทธิภาพ ระบบจะทิ้ง HLO ที่เพิ่มประสิทธิภาพแล้วหากใช้แฟล็ก -xla_dump_to ไปยังไฟล์ที่มีคำต่อท้ายชื่อไฟล์ "after_optimizations.txt" หากต้องการทิ้ง HLO หลังจากพาสระดับกลางที่เปลี่ยน HloModule จริงๆ คุณสามารถใช้แฟล็ก -xla_dump_hlo_pass_re=.* (หรือนิพจน์ทั่วไปที่เฉพาะเจาะจงเพื่อจำกัดเฉพาะพาสบางรายการ)

Scheduling

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

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

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

การกำหนดบัฟเฟอร์

ก่อนที่จะลดระดับเป็น LLVM IR เราจะเรียกใช้การส่งผ่านการกำหนดบัฟเฟอร์ที่จะกำหนดส่วนบัฟเฟอร์ให้กับแต่ละคำสั่งในกราฟ HLO การกำหนดบัฟเฟอร์ จะดำเนินการในหลายขั้นตอน ดังนี้

  1. HloDataflowAnalysisกำหนด HloValues (บัฟเฟอร์เชิงตรรกะโดยพื้นฐาน) ให้กับ คำสั่ง สำหรับการดำเนินการในตำแหน่ง คุณสามารถนำ HloValue ของตัวถูกดำเนินการกลับมาใช้ใหม่ได้ Op อาจกำหนด HloValue มากกว่า 1 รายการ (เช่น มีรูปร่างผลลัพธ์เป็น Tuple)

  2. HloAliasAnalysis พยายามรวมบัฟเฟอร์สำหรับการดำเนินการแทนชื่อ และ คำนวณการแมปจาก HloValue ไปยัง HloBuffer

  3. BufferAssignment จะคำนวณการแมปของ HloBuffers กับส่วนบัฟเฟอร์ภายในบัฟเฟอร์ขนาดใหญ่ ในลักษณะที่ส่วนบัฟเฟอร์เดียวกันจะไม่ถูกใช้สำหรับ HloBuffers ที่แตกต่างกัน ซึ่งมีช่วงเวลาที่ซ้อนทับกัน สำหรับ Ops ที่อาจมีนามแฝง คุณสามารถมีช่วงเวลาที่ทับซ้อนกันเล็กน้อยได้ (เวลาสิ้นสุดของรายการหนึ่ง HloBuffer อาจตรงกับเวลาเริ่มต้นของอีกรายการหนึ่ง HloBuffer) เมื่อใช้แฟล็ก -xla_dump_to ระบบจะทิ้งข้อมูลบางอย่างเกี่ยวกับการกำหนดบัฟเฟอร์ลงในไฟล์ที่มีคำต่อท้ายชื่อว่า "after_optimizations-buffer-assignment.txt"

Thunks

หลังจากเพิ่มประสิทธิภาพและกำหนดเวลาของกราฟ HLO แล้ว จะมีการลดระดับกราฟลงเป็น ลำดับเชิงเส้นของ Thunk สำหรับแบ็กเอนด์ที่เฉพาะเจาะจง (CPU หรือ GPU)

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

การปล่อยก๊าซที่เกิดจากความคิด

กระบวนการแปลงการคำนวณ HLO ที่กำหนดเวลาไว้เป็นลำดับ Thunk เรียกว่า "การปล่อย Thunk" ซึ่งจะได้รับการจัดการโดยคลาสการปล่อยเฉพาะในแต่ละแบ็กเอนด์

สำหรับแบ็กเอนด์ GPU จะมีการจัดการโดย IrEmitterUnnested EmitHloComputation จะวนซ้ำในรายการที่กำหนดเวลาไว้ของคำสั่ง HLO ใน การคำนวณและส่งไปยังเมธอด Emit... ที่เชี่ยวชาญ (เช่น EmitFusion, EmitConvolutionThunk, EmitWhile) แต่ละวิธีเหล่านี้จะสร้างออบเจ็กต์ Thunk ที่เหมาะสมและต่อท้ายออบเจ็กต์เหล่านั้นในลำดับ Thunk

สำหรับแบ็กเอนด์ CPU ThunkEmitter จะทำหน้าที่นี้และจัดระเบียบในลักษณะที่คล้ายกัน ThunkSequence จะฝังอยู่ใน CpuExecutable

โปรดทราบว่าแต่ละคำสั่งในการคำนวณรายการของโมดูล HLO อาจ สอดคล้องกับ Thunk 0 รายการ (kTuple, kConstant, ..), 1 รายการ หรือหลายรายการ (เช่น คำสั่งเรียง) ในลำดับ Thunk สุดท้าย

บัฟเฟอร์คำสั่ง: การเพิ่มประสิทธิภาพการดำเนินการใน GPU

ฮาร์ดแวร์ GPU สมัยใหม่ช่วยให้บันทึกลำดับการทำงานของ GPU (การเปิดใช้เคอร์เนล การคัดลอกหน่วยความจำ ฯลฯ) ได้ครั้งเดียว แล้วเล่นลำดับซ้ำหลายครั้ง โดยมีค่าใช้จ่ายของ CPU น้อยที่สุด การเพิ่มประสิทธิภาพนี้มีความสำคัญอย่างยิ่ง โดยเฉพาะอย่างยิ่งสำหรับภาระงานที่มีเคอร์เนลขนาดเล็กจำนวนมากที่เปิดตัวอย่างรวดเร็ว XLA ใช้Command Buffer เป็นการแยกข้อมูลกราฟ CUDA หรือกราฟ HIP อินเทอร์เฟซหลัก กำหนดไว้ใน GpuCommandBuffer

บัฟเฟอร์คำสั่งจะแสดงในลำดับ Thunk โดย CommandBufferThunk

Emitter ไม่ได้สร้าง Thunk นี้จากคำสั่ง HLO โดยตรง แต่จะดำเนินการโดย CommandBufferConversionPass ที่ทำงานใน ThunkSequence เอง

การส่งผ่านจะระบุลำดับย่อยที่ต่อเนื่องกันของ Thunk ที่เข้ากันได้ (เช่น ชุดของ KernelThunk และ GemmThunk) จากนั้นจะแทนที่ลำดับย่อยที่พบด้วย CommandBufferThunk เดียว Thunk ใหม่จะห่อหุ้มตรรกะของ Thunk เดิมเป็นรายการออบเจ็กต์ CommandBufferCmd แบบเบา เมื่อ CommandBufferThunk ทำงานเป็นครั้งแรกในสตรีม GPU ที่กำหนด จะ "บันทึก" ลำดับคำสั่งลงในบัฟเฟอร์คำสั่งฮาร์ดแวร์ ในการดำเนินการครั้งต่อๆ ไปทั้งหมด ระบบจะออกคำสั่งเดียวไปยัง GPU เพื่อ "เล่นซ้ำ" ลำดับที่บันทึกไว้ ซึ่งจะช่วยหลีกเลี่ยงค่าใช้จ่ายของ CPU ในการเปิดตัวเคอร์เนลแต่ละรายการ

ไฟล์ปฏิบัติการ

ผลิตภัณฑ์สุดท้ายของไปป์ไลน์การคอมไพล์ XLA คือไฟล์ที่เรียกใช้งานได้แบบแพลตฟอร์มเฉพาะ ที่ทำงานได้ด้วยตัวเอง ออบเจ็กต์นี้จะห่อหุ้มข้อมูลทั้งหมดที่จำเป็นในการเรียกใช้โปรแกรมที่คอมไพล์แล้ว ในอุปกรณ์เป้าหมาย เช่น CPU หรือ GPU ซึ่งเป็นตัวเชื่อมระหว่างคอมไพเลอร์ กับรันไทม์ รันไทม์สมัยใหม่ เช่น PJRT จะใช้การแยกส่วนระดับที่สูงขึ้นเล็กน้อย (ดู PjRtExecutable) แต่ท้ายที่สุดแล้วสิ่งเหล่านี้จะห่อหุ้มไฟล์ที่เรียกใช้งานได้เฉพาะแบ็กเอนด์

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

  • โค้ดที่คอมไพล์แล้ว: นี่คือโค้ดเครื่องระดับต่ำที่จะทำงานบนอุปกรณ์ สำหรับ CPU โดยปกติแล้วจะเป็นไฟล์ออบเจ็กต์อย่างน้อย 1 ไฟล์ สำหรับ GPU นี่คือ โค้ดอุปกรณ์ที่คอมไพล์แล้วในรูปแบบ PTX หรือ HSACO ซึ่งโหลดลงใน GPU ที่ รันไทม์

  • แผนการดำเนินการ (ThunkSequence): หัวใจสำคัญของตรรกะรันไทม์ นี่คือ ลำดับเชิงเส้นของออบเจ็กต์ Thunk แต่ละ Thunk แสดงถึงหน่วยงานเดียว เช่น การเปิดตัวเคอร์เนล การเรียกใช้ฟังก์ชันไลบรารี (เช่น cuBLAS) หรือ การจัดการโฟลว์การควบคุม รันไทม์จะเรียกใช้โปรแกรมโดยการวนซ้ำในลำดับนี้

  • เลย์เอาต์หน่วยความจำ (BufferAssignment): ข้อมูลเมตาที่สำคัญนี้สร้างขึ้นโดย BufferAssigner และอธิบายเลย์เอาต์หน่วยความจำทั้งหมดสำหรับการคำนวณ โดยจะระบุขนาดของบัฟเฟอร์ทั้งหมด รวมถึงวิธีจัดสรรและนำหน่วยความจำกลับมาใช้ใหม่ สำหรับพารามิเตอร์ เอาต์พุต และค่าชั่วคราว รันไทม์จะใช้ข้อมูลนี้เพื่อ จัดสรรหน่วยความจำของอุปกรณ์และส่งพอยน์เตอร์ที่ถูกต้องไปยังแต่ละ Thunk

  • (ไม่บังคับ) โมดูล HLO: สำหรับการแก้ไขข้อบกพร่องและการจัดทำโปรไฟล์ Executable มักจะ เก็บการอ้างอิงไปยัง HloModule สุดท้ายที่ได้รับการเพิ่มประสิทธิภาพซึ่งคอมไพล์ จาก

คอมไพเลอร์จะประสานงานการสร้างไฟล์ที่เรียกใช้งานได้สุดท้ายสำหรับแต่ละ แบ็กเอนด์ที่เฉพาะเจาะจง RunBackendเมธอดของการติดตั้งใช้งานคอมไพเลอร์เป็นขั้นตอนสุดท้ายในกระบวนการคอมไพล์ ซึ่งจะแพ็กเกจอาร์ติแฟกต์ที่คอมไพล์ทั้งหมด เป็นออบเจ็กต์ที่เรียกใช้งานได้ GpuCompiler และ CpuCompiler กำหนดเป้าหมายเป็น GPU และ CPU ตามลำดับ

เมื่อผู้ใช้เรียกใช้ Execute... ในไฟล์ที่เรียกใช้งานได้ รันไทม์จะใช้ BufferAssignment เพื่อจัดสรรหน่วยความจำ จากนั้นจะเรียกใช้ ThunkSequence เพื่อ เปิดตัวการดำเนินการในอุปกรณ์โดยใช้โค้ดที่คอมไพล์แล้ว