XLA (Accelerated Linear Algebra) คือคอมไพเลอร์แมชชีนเลิร์นนิง (ML) ที่เพิ่มประสิทธิภาพพีชคณิตเชิงเส้น โดยให้การปรับปรุงความเร็วในการดำเนินการและการใช้หน่วยความจำ หน้านี้แสดงภาพรวมโดยย่อเกี่ยวกับวัตถุประสงค์และสถาปัตยกรรมของคอมไพเลอร์ XLA
วัตถุประสงค์
ปัจจุบัน XLA รองรับส่วนหน้าของเฟรมเวิร์ก ML หลายตัว (รวมถึง PyTorch, TensorFlow และ JAX) และเป็นส่วนหนึ่งของโครงการ OpenXLA ซึ่งเป็นระบบนิเวศของเทคโนโลยีคอมไพเลอร์โอเพ่นซอร์สสำหรับ ML ที่ได้รับการพัฒนาร่วมกันโดยองค์กรฮาร์ดแวร์และซอฟต์แวร์ ML ชั้นนำ ก่อนที่จะสร้างโครงการ OpenXLA XLA ได้รับการพัฒนาภายในโครงการ TensorFlow แต่วัตถุประสงค์พื้นฐานยังคงเหมือนเดิม:
ปรับปรุงความเร็วในการดำเนินการ คอมไพล์กราฟย่อยเพื่อลดเวลาดำเนินการของ ops ที่มีอายุสั้น และกำจัดโอเวอร์เฮดจากรันไทม์ ฟิวส์การดำเนินการไปป์ไลน์เพื่อลดโอเวอร์เฮดของหน่วยความจำ และเชี่ยวชาญรูปร่างเทนเซอร์ที่รู้จักเพื่อให้สามารถเผยแพร่อย่างต่อเนื่องเชิงรุกมากขึ้น
ปรับปรุงการใช้หน่วยความจำ วิเคราะห์และกำหนดเวลาการใช้หน่วยความจำ กำจัดบัฟเฟอร์การจัดเก็บข้อมูลระดับกลางจำนวนมาก
ลดการพึ่งพาการดำเนินการแบบกำหนดเอง ขจัดความจำเป็นใน ops แบบกำหนดเองจำนวนมากโดยการปรับปรุงประสิทธิภาพของ ops ระดับต่ำที่หลอมรวมโดยอัตโนมัติ เพื่อให้ตรงกับประสิทธิภาพของ 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 รองรับ NVIDIA GPU ผ่านทางแบ็กเอนด์ LLVM NVPTX CPU backend รองรับ CPU ISA หลายตัว