การพัฒนาแบ็กเอนด์ใหม่สำหรับ XLA

คู่มือนี้มีไว้สำหรับวิศวกรระบบที่ต้องการใช้ XLA สร้างเอาต์พุตโปรแกรมที่กำหนดเป้าหมายฮาร์ดแวร์ของตนอย่างมีประสิทธิภาพ คู่มือนี้ไม่มีคำแนะนำให้ทีละขั้นตอนและจะถือว่ามีความรู้เกี่ยวกับ LLVM, Bazel และ XLA

XLA มีอินเทอร์เฟซสมมติที่สถาปัตยกรรมหรือ Accelerator ใหม่สามารถนำมาใช้เพื่อสร้างแบ็กเอนด์เพื่อเรียกใช้เอาต์พุตโปรแกรม ML โดย XLA ได้ การกำหนดเป้าหมาย XLA ใหม่น่าจะง่ายกว่าและปรับขนาดได้มากกว่าการใช้ตัวเลือกที่มีอยู่ทั้งหมดจากเฟรมเวิร์กฟรอนท์เอนด์ เช่น PyTorch หรือ TensorFlow สำหรับฮาร์ดแวร์ใหม่

การติดตั้งใช้งานส่วนใหญ่จะเป็นอย่างใดอย่างหนึ่งต่อไปนี้

  1. สถาปัตยกรรม CPU ที่มีอยู่ยังไม่ได้รับการสนับสนุนอย่างเป็นทางการจาก XLA โดยมีหรือไม่มีแบ็กเอนด์ LLVM ที่มีอยู่
  2. ฮาร์ดแวร์ที่ไม่ใช่ CPU ที่มีแบ็กเอนด์ LLVM อยู่แล้ว
  3. ฮาร์ดแวร์ที่ไม่ใช่ CPU ซึ่งไม่มีแบ็กเอนด์ LLVM ที่มีอยู่

สถานการณ์ที่ 1: XLA ยังไม่รองรับสถาปัตยกรรม CPU ที่มีอยู่อย่างเป็นทางการ

ในกรณีนี้ ให้เริ่มต้นด้วยการดูแบ็กเอนด์ CPU ของ XLA ที่มีอยู่ XLA ช่วยให้กำหนดเป้าหมาย CPU ต่างๆ ได้อย่างง่ายดายโดยใช้ LLVM เนื่องจากความแตกต่างหลักระหว่างแบ็กเอนด์ XLA สำหรับ CPU คือโค้ดที่สร้างโดย LLVM

หากผู้ให้บริการฮาร์ดแวร์มีแบ็กเอนด์ LLVM สำหรับฮาร์ดแวร์ ก็สามารถลิงก์แบ็กเอนด์กับ LLVM ที่สร้างด้วย XLA ได้อย่างง่ายดาย ในโหมด JIT แบ็กเอนด์ CPU ของ XLA จะปล่อยโค้ดสำหรับ CPU ของโฮสต์ สำหรับการคอมไพล์ล่วงหน้า xla::AotCompilationOptions สามารถเพิ่ม LLVM เพื่อกำหนดค่าสถาปัตยกรรมเป้าหมายได้ 3 เท่า

หากไม่มีแบ็กเอนด์ LLVM อยู่ แต่มีโปรแกรมสร้างโค้ดประเภทอื่นอยู่แล้ว คุณน่าจะนำแบ็กเอนด์ CPU ที่มีอยู่ส่วนใหญ่มาใช้ซ้ำได้

สถานการณ์ที่ 2: ฮาร์ดแวร์ที่ไม่ใช่ CPU ที่มีแบ็กเอนด์ LLVM อยู่แล้ว

เป็นไปได้ที่จะทำโมเดลการติดตั้งใช้งาน xla::Compiler ใหม่บนคลาส xla::CPUCompiler และ xla::GPUCompiler ที่มีอยู่ เนื่องจากคลาสเหล่านี้ปล่อย LLVM IR อยู่แล้ว จึงอาจเป็นไปได้ว่าจะต้องมีการเปลี่ยนแปลงหลายๆ ด้านในการสร้าง LLVM IR แต่สามารถแชร์โค้ดจำนวนมากกับแบ็กเอนด์ที่มีอยู่ได้

ตัวอย่างที่ดีที่ควรทำตามคือ แบ็กเอนด์ GPU ของ XLA แบ็กเอนด์ของ GPU จะกำหนดเป้าหมายเป็น ISA ที่ไม่ใช่ CPU ดังนั้นการสร้างโค้ดบางอย่างจึงแตกต่างกันไปตามโดเมน GPU ฮาร์ดแวร์ประเภทอื่นๆ เช่น DSP อย่าง Hexagon (ซึ่งมีแบ็กเอนด์ LLVM แบบอัปสตรีม) สามารถใช้ลอจิกการปล่อยก๊าซ LLVM IR ซ้ำได้ แต่ส่วนอื่นๆ จะแตกต่างออกไป

สถานการณ์ที่ 3: ฮาร์ดแวร์ที่ไม่ใช่ CPU ซึ่งไม่มีแบ็กเอนด์ LLVM ที่มีอยู่

หากใช้ LLVM ไม่ได้ ตัวเลือกที่ดีที่สุดคือใช้แบ็กเอนด์ใหม่สำหรับ XLA สำหรับฮาร์ดแวร์ที่ต้องการ ตัวเลือกนี้ต้องใช้ความพยายามมากที่สุด คลาสที่ต้องดำเนินการมีดังนี้

  • StreamExecutor: อุปกรณ์ส่วนใหญ่ไม่จำเป็นต้องใช้ StreamExecutor บางวิธี ดูรายละเอียดเกี่ยวกับการติดตั้งใช้งาน StreamExecutor ที่มีอยู่
  • xla::Compiler: คลาสนี้สรุปการคอมไพล์ HLO เข้าเป็น xla::Executable
  • xla::Executable: คลาสนี้ใช้เพื่อเริ่มการคำนวณแบบคอมไพล์บนแพลตฟอร์ม
  • xla::TransferManager: คลาสนี้ช่วยให้แบ็กเอนด์มอบกลไกเฉพาะแพลตฟอร์มสำหรับ การสร้างข้อมูลลิเทอรัล XLA จากแฮนเดิลหน่วยความจำของอุปกรณ์ที่กำหนด กล่าวอีกนัยหนึ่งคือ เครื่องมือนี้ช่วยห่อหุ้มการโอนข้อมูลจากโฮสต์ไปยังอุปกรณ์และย้อนกลับ