สถาปัตยกรรม XLA

XLA (Accelerated Linear Algebra) คือคอมไพเลอร์แมชชีนเลิร์นนิง (ML) ที่เพิ่มประสิทธิภาพพีชคณิตเชิงเส้นโดยปรับปรุงความเร็วในการดำเนินการและการใช้หน่วยความจำ หน้านี้ให้ภาพรวมคร่าวๆ เกี่ยวกับวัตถุประสงค์และสถาปัตยกรรมของคอมไพเลอร์ XLA

วัตถุประสงค์

ปัจจุบัน XLA รองรับฟรอนท์เอนด์ของเฟรมเวิร์ก ML หลายรายการ (รวมถึง PyTorch, TensorFlow และ JAX) และเป็นส่วนหนึ่งของโครงการ OpenXLA ซึ่งเป็นระบบนิเวศของเทคโนโลยีคอมไพเลอร์โอเพนซอร์สสำหรับ ML ที่พัฒนาร่วมกันโดยองค์กรชั้นนำด้านฮาร์ดแวร์และซอฟต์แวร์ ML ก่อนที่จะสร้างโปรเจ็กต์ OpenXLA มีการพัฒนาขึ้น XLA ภายในโปรเจ็กต์ TensorFlow แต่วัตถุประสงค์พื้นฐานยังคงเหมือนเดิม

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

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

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

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

วิธีการทำงาน

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

โดยทั่วไป กระบวนการคอมไพล์ที่แปลงกราฟโมเดลเป็นไฟล์ปฏิบัติการที่เพิ่มประสิทธิภาพตามเป้าหมายจะมีขั้นตอนเหล่านี้

  1. XLA ทำการเพิ่มประสิทธิภาพและการวิเคราะห์ในตัวหลายรายการบนกราฟ SttableHLO ที่ไม่ขึ้นอยู่กับเป้าหมาย เช่น CSE, ฟิวชันการดำเนินการที่ไม่ขึ้นกับเป้าหมาย และการวิเคราะห์บัฟเฟอร์สำหรับการจัดสรรหน่วยความจำรันไทม์สำหรับการคำนวณ ในช่วงการเพิ่มประสิทธิภาพนี้ XLA จะแปลงภาษา StableHLO เป็นภาษา HLO ภายในด้วย

  2. XLA จะส่งการคำนวณ HLO ไปยังแบ็กเอนด์เพื่อเพิ่มประสิทธิภาพระดับ HLO เพิ่มเติม โดยในครั้งนี้จะคำนึงถึงข้อมูลและความต้องการเฉพาะเป้าหมาย ตัวอย่างเช่น แบ็กเอนด์ของ GPU อาจทำการฟิวชันการดำเนินการซึ่งมีประโยชน์โดยเฉพาะสำหรับโมเดลการเขียนโปรแกรม GPU และกำหนดวิธีแบ่งพาร์ติชันการประมวลผลในสตรีม ในขั้นตอนนี้ แบ็กเอนด์อาจจับคู่รูปแบบหรือผสมการดำเนินการบางอย่างกับการเรียกไลบรารีที่เพิ่มประสิทธิภาพ

  3. จากนั้นแบ็กเอนด์จะสร้างโค้ดเฉพาะเป้าหมาย แบ็กเอนด์ CPU และ GPU ที่รวมอยู่ใน XLA จะใช้ LLVM สำหรับ IR, การเพิ่มประสิทธิภาพ และการสร้างโค้ดในระดับต่ำ แบ็กเอนด์เหล่านี้ปล่อย LLVM IR ที่จำเป็นในการแสดงการคำนวณ HLO ด้วยวิธีที่มีประสิทธิภาพ จากนั้นเรียกใช้ LLVM เพื่อปล่อยโค้ดแบบเนทีฟจาก LLVM IR นี้

ภายในกระบวนการนี้ คอมไพเลอร์ XLA จะแยกเป็นส่วนๆ ซึ่งทำให้สามารถนำไปใส่ในแบ็กเอนด์สำรองเพื่อกำหนดเป้าหมายสถาปัตยกรรม HW ใหม่บางรายการได้อย่างง่ายดาย ปัจจุบันแบ็กเอนด์ GPU รองรับ GPU ของ NVIDIA ผ่านแบ็กเอนด์ LLVM NVPTX แบ็กเอนด์ของ CPU รองรับ ISA ของ CPU หลายรายการ