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

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

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

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

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

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

ในสถานการณ์นี้ ให้เริ่มต้นด้วยการดูแบ็กเอนด์ของ 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 จากแฮนเดิลหน่วยความจำอุปกรณ์ที่กำหนด กล่าวอีกนัยหนึ่งคือ วิธีนี้ช่วยห่อหุ้มการโอนข้อมูลจากโฮสต์ไปยังอุปกรณ์และส่งกลับ