หน้านี้ได้รับการแปลโดย Cloud Translation API
Switch to English

MLIR CodeGen สำหรับ XLA

XLA ทำงานบน HloInstruction และทำการปรับแต่งมากมายสำหรับการแสดงนี้โดยแบ่งปันสิ่งเหล่านี้จำนวนมากระหว่างอุปกรณ์เป้าหมาย ในบางจุดจะมีการคำนวณกำหนดการเชิงเส้นและบัฟเฟอร์หน่วยความจำถูกกำหนดให้กับแต่ละค่าแบบสแตติก codegen เฉพาะของอุปกรณ์ทำงานโดยการข้ามลำดับนี้และเรียก "emitters" เพื่อสร้างการแสดงที่เหมาะสมสำหรับอุปกรณ์ (เช่นฟังก์ชัน LLVM เดียวต่อการคำนวณ XLA บน CPU หรือลำดับของ "thunks" encapsulating การทำงานของ GPU และอาจสร้าง PTX เมื่อ กำหนดเป้าหมาย GPU)

ในฐานะขั้นตอนการจัดเตรียมเรากำลังอยู่ในขั้นตอนการดักจับกระบวนการหลังจาก XLA เสร็จสิ้นขั้นตอนการกำหนดบัฟเฟอร์และปล่อยแทนโมดูล MLIR ในภาษา lhlo จากนั้นเราทำการ codegen โดยใช้ส่วนประกอบ MLIR (ส่วนใหญ่ Linalg, affine และ GPU ส่วนใหญ่) ขึ้นอยู่กับอุปกรณ์

ด้านล่างนี้เป็นแผนของการบันทึกเพื่อย้าย XLA / GPU แบบเพิ่มหน่วยโดยใช้ lhlo เป็นอินพุต codegen

งาน

เจ้าภาพ เครื่อง
รูปแบบอินพุต HloInstruction * (ภารกิจ 1) HloInstruction * (ภารกิจ 1)
รูปแบบผลลัพธ์ xla :: Thunk (ภารกิจ 2) LLVM IR (ภารกิจ 3)
  • ภารกิจที่ 1 เปลี่ยนทั้งโฮสต์และรูปแบบอินพุตอุปกรณ์จาก HloInstruction * เป็น LHLO
  • ภารกิจที่ 2 เปลี่ยนรูปแบบเอาต์พุตของโฮสต์จาก thunks เป็น "landing pad สำหรับ host" (ดูด้านล่าง)
  • ภารกิจที่ 3 โอนย้ายอุปกรณ์เอาต์พุตจาก LLVM IR เป็น MLIR บางรูปแบบ เป็นตัวเลือกสำหรับโครงการนี้และดูส่วน "การโอนย้ายอุปกรณ์ LLVM IR" สำหรับรายละเอียด

โครงการนี้ให้ความสำคัญกับการมีโมเดลที่รันได้แบบ end-to-end โดยเปิดใช้งานตัวปล่อย LHLO ให้มากที่สุด นี่ก็หมายความว่ารายการลำดับวัตถุประสงค์ต่อไปนี้ตามลำดับความสำคัญ:

  • ทำให้ XLA / GPU ทำงานได้ด้วยตัวปล่อย LHLO โดยมี Thunks และ emitters ที่มีอยู่โดยไม่ได้แก้ไข
  • กำจัดการอ้างอิงถึง HloInstruction * ใน LHLO เป็นกรณี ๆ :
    • สลับอีซีแอลที่สืบทอดมาเป็นอีซีแอลที่ใช้ MLIR (เช่น Linalg) หรือ
    • แปลกลไก emitter ที่มีอยู่เพื่อใช้แทน MLIR (โอนย้ายเป็นมาตรฐานด้วยภาษาถิ่นของ GPU)

การย้าย Thunks (ภารกิจ 2)

xla :: gpu :: Thunk เป็นโครงสร้างข้อมูลที่:

  • สามารถเรียกใช้จากโฮสต์ (xla :: gpu :: Thunk :: ExecuteOnStream ())
  • ดำเนินการข้อมูลต่าง ๆ ในคลาสย่อย
  • โต้ตอบกับ BufferAllocation :: Slice และ StreamExecutor
  • เปิดตัวเมล็ดข้าว
  • โทรเข้าสู่ไลบรารีรันไทม์ทั้งหมด

ค่าใช้จ่ายที่รวมถึง:

  • การแทนข้อมูลการกำหนดค่า op-specific (เช่น convolution configs)
  • การโอนย้ายรูปร่าง op และรูปร่าง operand
  • เป็นตัวแทนของต้นไม้ thunks (ในขณะที่สภาพ ฯลฯ )

งานการโยกย้ายไม่ขึ้นอยู่กับการโยกย้าย LHLO / emitter ภายใต้ทรัพยากรที่ จำกัด มีการจัดลำดับความสำคัญหลังการโยกย้าย LHLO / อีซีแอล

เรามีหลายทางเลือกในการลดส่วนของโฮสต์จาก LHLO:

  • TFRT
    • (Pro) ชุด CUDA และ HIP ที่ยอดเยี่ยมสำหรับการใช้งาน
    • (Pro) ง่ายต่อการใช้งานการเรียกใช้ไลบรารี (cuDNN, cuBLAS, cuFFT, ฯลฯ ) เนื่องจากการแปลรหัส TFRT จะถูกตีความด้วยรหัส C ++
    • (Con) ด้านโฮสต์อยู่ระหว่างการพัฒนาและไม่ได้ทดสอบ
  • รหัส CPU Jitted
    • (Pro) ความสามารถต่ำกว่ามาก สร้างลูปและเงื่อนไขสองสามอย่างแล้วเสร็จ
    • (Con) GPUDialect ยังไม่ได้จำลองการจัดสรรเชน / สตรีม / asynchronicity / อุปกรณ์
    • (Con) การสนับสนุนรันไทม์ CUDA / HIP นั้นน้อยมาก (เส้นทางชุดเครื่องมือรุ่นการโหลดแบบไดนามิก ฯลฯ )
  • รันไทม์ XLA ที่มีอยู่ (ตีความ)

การตัดสินใจ: ใช้ TFRT แต่ยังสนับสนุนการ jitting CPU code ใน TFRT

การย้ายอุปกรณ์ LLVM IR (ภารกิจ 3)

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

ElementalIrEmitter มีลักษณะเฉพาะใน:

  • ส่วนใหญ่ของรหัสที่ใช้ร่วมกันระหว่าง XLA / GPU และ CPU
  • มันแสดงถึงส่วนใหญ่ของ ops ที่เห็นในโมเดลรวมถึง ops ที่ชาญฉลาดทั้งหมด
  • การหลอมรวมส่วนใหญ่ขึ้นอยู่กับ ElementalIrEmitter
  • มันเป็นโครงสร้างที่เรียบง่ายเนื่องจากอธิบาย DAG การพึ่งพาข้อมูลระหว่างองค์ประกอบ op และองค์ประกอบ operand
  • ส่วนใหญ่เป็นแบบพกพาและระดับสูง (เช่นแตกต่างจาก GPU kReduce และ GPU kCopy)
  • การสนับสนุนรูปร่างแบบไดนามิกนั้นง่ายสำหรับ ops ที่ชาญฉลาดอย่างน้อยองค์ประกอบ

ตอนนี้สำหรับ ops ทั้งหมดที่ปล่อยออกมาเป็นองค์ประกอบหรือไม่มีหลายสถานะของสถานะสิ้นสุดของแต่ละ XLA op:

  1. รหัสอุปกรณ์ยังคงเป็น LLVM IR
  2. Refactor ตัวปล่อยเดิมให้เป็นเหมือน LHLO -> MLIR LLVM สำเนียง:
    • (ราคา) จะถูกทิ้งหากเราต้องการย้ายไปยัง Standard ในที่สุด
    • (ประโยชน์) มันง่ายและมีกลไก สามารถทำได้ในระยะเวลาอันสั้น
    • (ประโยชน์) มันไม่ได้ประโยชน์มากกว่าเมื่อเปรียบเทียบกับ (1)
  3. Refactor ตัวส่งสัญญาณเก่าเป็นเหมือน LHLO -> MLIR GPU + มาตรฐาน + ลูป:
    • (ค่าใช้จ่าย) การยกตัวปล่อยที่มีอยู่ให้เป็นมาตรฐานเป็นการท้าทายบางอย่าง พอยน์เตอร์และ GEP จำเป็นต้องแปลงเป็น MemRefs และ SubViews การทำให้มั่นใจว่า amdgpu สมบูรณ์นั้นเป็นอีกสิ่งหนึ่ง
    • (ราคา) XLA / GPU ใช้ข้อมูลเมตา LLVM อย่างมาก:
      • range สำหรับดัชนีบล็อก / เธรด
      • align dereferenceable , invariant.load , alias.scope , noalias สำหรับโหลด / ร้านค้า
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable สำหรับการวนซ้ำตามลำดับ
    • (ประโยชน์) สามารถระยะยาว พกพาได้มากกว่า
  4. Refactor ตัวปล่อยเก่าให้เป็น LHLO -> Linalg และเขียนตัวปล่อย Linalg ใหม่
    • (ราคา) กรณีนี้เป็นกรณี ๆ ไป เปรียบเทียบกับตัวเลือกก่อนหน้าการใช้งานใหม่ที่ตรงกับประสิทธิภาพของ XLA ต้องผ่านเกณฑ์มาตรฐาน <-> ปรับขั้นตอนการทำงานให้เหมาะสมซึ่งอาจเป็นค่าใช้จ่ายที่สำคัญสำหรับบาง ops
    • (ประโยชน์) สแต็คแบบรวม การสนับสนุนจากชุมชน พกพา; ศักยภาพการเพิ่มประสิทธิภาพมากขึ้น

สรุป:

  • อย่าไปเพื่อ (2) (1) หรือ (3) ดีกว่า (2) (2) ค่าใช้จ่ายมากกว่า (1) เนื่องจากต้องมีการปรับโครงสร้างเชิงกลจำนวนมาก ด้วย (1) เรายังคงสามารถบรรลุเป้าหมายของการเปิดใช้งาน XLA เพื่อรับตัวปล่อย MLIR นี่คือการทำ LHLO -> LLVM IR -> เรียกใช้อุปกรณ์ส่งสัญญาณรุ่นเก่า
  • ElementalIrEmitter ops ใช้สำหรับ (4) แต่ไม่เพิ่มขึ้น ไม่มีทางที่จะทำ op โดย op เพราะ ops ที่ปล่อยออกมาเป็นองค์ประกอบทั้งหมดจะเชื่อมต่อกับกราฟเดียวกัน งานนี้ยังสามารถทำหน้าที่เป็นจุดรวมพลังของกองกำลังที่กำลังดำเนินการอยู่หลายแห่ง (xla / service / mlir_gpu, ตัวสร้างเคอร์เนล, Linalg)
  • ตัวเลือกอื่น ๆ ทั้งหมดใช้สำหรับ (1) เป้าหมายอาจยืดเยื้อพวกเขาอาจถูกย้ายไปที่ (3) หรือ (4)

จัดลำดับความสำคัญ

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

การจัดลำดับความสำคัญคือ: Task1 (LHLO สำหรับผู้ปล่อยมรดก)> ภารกิจ 2 (Thunks)> ภารกิจ 3 (MLIR emitters)

ในตอนท้ายของภารกิจที่ 1 ผู้ใช้ XLA สามารถสร้าง LHLO (เช่นตัวสร้างเคอร์เนล) และดำเนินการได้ รูปแบบการรวบรวมจะไม่สามารถทำให้เป็นลำดับ MLIR

ในตอนท้ายของภารกิจที่ 2 LHLO จะลดค่า MLIR ที่เหมาะสมและต่อเนื่องได้ สิ่งนี้ทำให้การรวบรวมแบบออฟไลน์

ในตอนท้ายของภารกิจที่ 3 ตัวปล่อย XLA ทั้งหมดนั้นใช้ MLIR ในการปรับใช้

ออกแบบรายละเอียด

ขั้นตอนที่ 1: (ภารกิจที่ 1) ทำ LHLO ให้สมบูรณ์และสร้างตัวปล่อยมรดกใช้ LHLO

ขั้นตอนนี้ทำให้ตัวปล่อย XLA / GPU ที่มีอยู่ทั้งหมดโต้ตอบกับ MLIR ops ขั้นตอนนี้เป็นการปรับโครงสร้างซ้ำอย่างสมบูรณ์และ NFC

ขั้นตอนนี้ส่วนใหญ่เป็นเชิงกล แต่ก็ควรสังเกตความแตกต่างดังต่อไปนี้ระหว่าง HloComputation ที่ไม่ได้ทำการทดสอบกับ LHLO:

  • HloInstruction แต่ละรายการมีการเข้าถึงตัวถูกดำเนินการโดยตรง (DAG การไหลของข้อมูล) ในทางตรงกันข้าม LHLO op แต่ละอันมีสิทธิ์เข้าถึงตัวถูกดำเนินการบัฟเฟอร์ (bipartite ระหว่าง ops และ buffer) LHLO ops ต้องผ่านการใช้โซ่ def เพื่อเข้าถึงตัวถูกดำเนินการ
  • มรดกที่ยังไม่ผ่านการปล่อย emitters สังเกตุแทบไม่เคยเข้าถึงตัวถูกดำเนินการของพวกเขา ข้อยกเว้นเพียงอย่างเดียวคือ kReduce
  • ตัวปล่อยมรดกที่ยังไม่ผ่านการเข้าถึงเข้าถึง BufferAssignment เพียงเพื่อรับส่วนข้อมูลไม่ใช่สำหรับการเข้าถึงโครงสร้างข้อมูล aux เช่น dataflow_analysis () หรือ alias_analysis () llvm_ir สร้าง alias_analysis () ของตัวเองตามข้อมูลส่วนแบ่ง

บทสรุปก็คือว่า LHLO ควรจะเข้าทางด้านขวาโดยไม่ต้องยุ่งยากใหญ่

ขั้นตอนที่ 2: (ตัวเลือก) การสนับสนุนการทำโปรไฟล์

ขั้นตอนนี้จำเป็นเฉพาะเมื่อเราเริ่มทิ้งตรรกะ XLA Thunk บางส่วน (ดูขั้นตอนถัดไป)

ก่อนที่จะเปิดตัวส่งสัญญาณที่ใช้ MLIR ใด ๆ เราจำเป็นต้องมีการทำโปรไฟล์สำหรับตัวปล่อยที่ใช้ MLIR

ปัจจุบัน XLA ดำเนินการทำโปรไฟล์ของตนเองโดยการเรียกใช้ตัวจับเวลาของ StreamExecutor ตัวจับเวลาภายใต้ประทุนจะแทรกสองเหตุการณ์ก่อนและหลังเคอร์เนลเรียกใช้และวัดเวลาการซิงค์ระหว่างสองเหตุการณ์

มีสามวิธีในการสนับสนุนการทำโปรไฟล์ใน MLIR:

  • เรียกใช้ profiler end-to-end
  • เพิ่มโปรไฟล์ op สำหรับแต่ละ op ใน LHLO โดยใช้ตัวสร้างโปรไฟล์

วิธีการแบบ end-to-end โปร่งใสกับ MLIR แต่ประสบปัญหาเดียวกันกับที่ทำให้ XLA ไม่ใช้ในตอนแรก: การเรียกไลบรารีที่รวบรวมโดย profiler (nvprof / ... ) ไม่สามารถเกี่ยวข้องกับ HLO ได้อย่างง่ายดาย Ops ตัวอย่างเช่น cuDNN เปิดใช้งานเคอร์เนลหลายอันสำหรับแต่ละ HLO และมันยากที่จะบอกว่าเมล็ดใดสอดคล้องกับ HLO ใด

วิธีการ "ผู้สร้างโปรไฟล์" ต้องการ:

  • LHLO ใช้ profiler เป็นพารามิเตอร์
  • การแทรก profile.start / profile.end ก่อนและหลังแต่ละ op
  • ผ่านจากโปรไฟล์ที่ลดลง {start, end} เพื่อการใช้งาน C ++

การทำโปรไฟล์ที่แน่นอนไม่สามารถทำได้อย่างง่ายดายสำหรับตัวเลือกที่สร้างโดย MLIR เนื่องจาก:

  • MLIR ไม่มีตัวจับเวลาและไม่ขึ้นอยู่กับ TFRT / StreamExecutor
  • MLIR ไม่สามารถเรียกใช้ฟังก์ชัน C ได้อย่างง่ายดายด้วยพารามิเตอร์ที่ซับซ้อน

ขั้นตอนที่ 3: (ภารกิจ 2) การโอนย้าย Thunks

ตามที่บันทึกไว้มี Thunks ประมาณสามชนิด:

  • KernelThunk ซึ่งเรียกใช้เคอร์เนล
  • thunks ควบคุมการไหลซึ่งมีลอจิกโฟลว์การควบคุมโฮสต์
  • คลังไลบรารี: cuDNN, cuBLAS, cuFFT, NCCL, เป็นต้น

แผนคือ:

  • ทำให้ Thunks (de) เป็นอนุกรม
  • ช่วยปรับปรุง TFRT ให้อยู่ในสถานะที่สามารถรองรับความหมายเหล่านี้
  • เมื่อสถานะดีขึ้นให้โยกย้ายส่วนบุคคลทีละส่วน

รายการการดำเนินการเหล่านี้มีการสั่งซื้อเพียงบางส่วนเท่านั้น คำสั่งการดำเนินการตามจริง / วิศวกรรมขนานจะถูกประเมินตามขั้นตอน

ขั้นตอนที่ 4: (งาน 3) โยกย้าย ElementalIrEmitter

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

โปรดสังเกตว่ามันเป็นประโยชน์ในการโยกย้าย ElementalIrEmitter ของ XLA / CPU เช่นกันเนื่องจากพวกเขาแบ่งปันรหัสส่วนใหญ่

เมื่อการเปรียบเทียบและการล่าสัตว์ดำเนินการเสร็จสิ้นแล้ว (TODO: กำหนดความเท่าเทียมกันของผลการปฏิบัติงาน) เราจะเปิดตัวปล่อยองค์ประกอบตาม MLIR ใหม่และลบ ElementalIrEmitter ดั้งเดิมออก

ขั้นตอนนี้ยังให้การเปลี่ยนฟิวชั่นได้ง่าย (ops ที่ซ้อนกัน) สำหรับการโอนย้ายในภายหลัง

ขั้นตอนที่ 5: รองรับ Multi-Stream หรือ Drop

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

ขั้นตอนที่ 6: (ภารกิจ 3) Ops อุปกรณ์ที่ย้ายข้อมูล

ขั้นตอนนี้โอนย้าย ops ที่ไม่ได้ใช้ทั้งหมดจากนั้นเราสามารถลบ emitters ที่ไม่ได้ใช้ทั้งหมด

สิ่งนี้เรียกใช้ rewrite / refactor สำหรับ kCopy และ kReduce kReduce ได้ทำงานมาแล้วมากมายดังนั้นจำนวนงานที่ต้องทำจริงยังคงปรากฏให้เห็น