XLA (Accelerated Linear Algebra) เป็นคอมไพเลอร์แมชชีนเลิร์นนิง (ML) ที่เพิ่มประสิทธิภาพพีชคณิตเชิงเส้นโดยปรับปรุงความเร็วในการดำเนินการและการใช้หน่วยความจำ หน้านี้จะแสดงภาพรวมโดยสรุปของวัตถุประสงค์และสถาปัตยกรรมของคอมไพเลอร์ XLA
วัตถุประสงค์
ปัจจุบัน XLA รองรับฟรอนท์เอนด์ของเฟรมเวิร์ก ML หลายรายการ (รวมถึง PyTorch, TensorFlow และ JAX) และเป็นส่วนหนึ่งของโครงการ OpenXLA ซึ่งเป็นระบบนิเวศของเทคโนโลยีคอมไพเลอร์โอเพนซอร์สสำหรับ ML ที่พัฒนาขึ้นร่วมกันโดยองค์กรชั้นนำด้านฮาร์ดแวร์และซอฟต์แวร์ ML ก่อนที่จะสร้างโปรเจ็กต์ OpenXLA มีการพัฒนาขึ้น XLA ภายในโปรเจ็กต์ TensorFlow แต่วัตถุประสงค์พื้นฐานจะยังคงเหมือนเดิม ดังนี้
ปรับปรุงความเร็วในการดำเนินการ คอมไพล์กราฟย่อยเพื่อลดเวลาในการดำเนินการที่ดำเนินไปเป็นระยะเวลาสั้นๆ และขจัดโอเวอร์เฮดออกจากรันไทม์ รวมการดำเนินการแบบไปป์ไลน์เพื่อลดปริมาณโอเวอร์เฮดของหน่วยความจำ และสร้างรูปแบบ tensor ที่รู้จักเพื่อให้สามารถใช้การเผยแพร่ได้อย่างสม่ำเสมอมากขึ้น
ปรับปรุงการใช้งานหน่วยความจำ วิเคราะห์และกำหนดเวลาการใช้หน่วยความจำโดยลดบัฟเฟอร์ของ พื้นที่เก็บข้อมูลระหว่างกลางจำนวนมาก
ลดการใช้การดำเนินการที่กำหนดเอง ขจัดความจำเป็นของการดำเนินการที่กำหนดเองจำนวนมากด้วยการปรับปรุงประสิทธิภาพของ Ops ระดับต่ำที่รวมเข้าด้วยกันโดยอัตโนมัติเพื่อให้เข้ากับประสิทธิภาพของการดำเนินการที่กำหนดเองซึ่งก่อนหน้านี้รวมเข้าด้วยกันด้วยมือ
ปรับปรุงการถ่ายโอนได้ ช่วยทำให้การเขียนแบ็กเอนด์ใหม่สำหรับฮาร์ดแวร์ใหม่เป็นเรื่องง่าย เพื่อให้โมเดล ML จำนวนมากทำงานได้โดยไม่มีการแก้ไขบนฮาร์ดแวร์นั้น ซึ่งตรงกันข้ามกับแนวทางการความเชี่ยวชาญแบบโมโนลิธแต่ละรายการสำหรับฮาร์ดแวร์ใหม่ ซึ่งกำหนดให้ต้องเขียนโมเดลใหม่เพื่อให้ใช้ประโยชน์จากการดำเนินการเหล่านั้นได้
วิธีการทำงาน
คอมไพเลอร์ XLA ใช้กราฟโมเดลจากเฟรมเวิร์ก ML ที่กำหนดไว้ใน StableHLO และคอมไพล์เป็นวิธีการของเครื่องสำหรับสถาปัตยกรรมต่างๆ StableHLO กำหนดชุดการดำเนินการที่มีเวอร์ชัน (HLO = การดำเนินการระดับสูง) ที่ให้เลเยอร์ความสามารถในการถ่ายโอนได้ระหว่างเฟรม ML และคอมไพเลอร์
โดยทั่วไป กระบวนการคอมไพล์ที่แปลงกราฟโมเดลเป็นไฟล์ปฏิบัติการที่เพิ่มประสิทธิภาพตามเป้าหมายจะรวมถึงขั้นตอนเหล่านี้
XLA จะดำเนินการส่งการเพิ่มประสิทธิภาพและการวิเคราะห์ในตัวหลายรายการบนกราฟ StableHLO ที่ไม่ขึ้นอยู่กับเป้าหมาย เช่น CSE ฟิวชันการดำเนินการที่ไม่ขึ้นกับเป้าหมาย และการวิเคราะห์บัฟเฟอร์สำหรับการจัดสรรหน่วยความจำรันไทม์สำหรับการคำนวณ ในระหว่างขั้นตอนการเพิ่มประสิทธิภาพนี้ XLA จะแปลงภาษา StableHLO เป็นภาษา HLO ภายในด้วย
XLA จะส่งการคำนวณ HLO ไปยังแบ็กเอนด์เพื่อเพิ่มประสิทธิภาพระดับ HLO เพิ่มเติม โดยครั้งนี้จะคำนึงถึงข้อมูลและความต้องการเฉพาะเป้าหมาย ตัวอย่างเช่น แบ็กเอนด์ของ GPU อาจดำเนินการฟิวชันการดำเนินการที่เป็นประโยชน์สำหรับโมเดลการเขียนโปรแกรม GPU โดยเฉพาะ และจะกำหนดวิธีการแบ่งพาร์ติชันการประมวลผลในสตรีม ในขั้นตอนนี้ แบ็กเอนด์อาจจับคู่รูปแบบหรือการดำเนินการบางรายการกับการเรียกใช้ไลบรารีที่เพิ่มประสิทธิภาพ
จากนั้นแบ็กเอนด์จะดำเนินการสร้างโค้ดเฉพาะเป้าหมาย แบ็กเอนด์ของ CPU และ GPU ที่มาพร้อมกับ XLA จะใช้ LLVM สำหรับ IR, การเพิ่มประสิทธิภาพ และการสร้างโค้ดระดับต่ำ แบ็กเอนด์เหล่านี้จะปล่อย LLVM IR ที่จำเป็นในการแสดงการคำนวณ HLO อย่างมีประสิทธิภาพ จากนั้นเรียกใช้ LLVM เพื่อปล่อยโค้ดเนทีฟจาก LLVM IR นี้
ในกระบวนการนี้ คอมไพเลอร์ XLA จะแยกเป็นโมดูลเพื่อให้นำไปใส่ในแบ็กเอนด์ทางเลือกสำหรับกำหนดเป้าหมายสถาปัตยกรรม HW แบบใหม่ๆ ได้ง่ายๆ ปัจจุบันแบ็กเอนด์ GPU รองรับ GPU NVIDIA ผ่านแบ็กเอนด์ LLVM NVPTX แบ็กเอนด์ของ CPU รองรับ ISA ของ CPU หลายรายการ