XLA (Accelerated Doğrusal Cebir), doğrusal cebiri optimize ederek yürütme hızı ve bellek kullanımında iyileştirmeler sağlayan bir makine öğrenimi (ML) derleyicisidir. Bu sayfada, XLA derleyicisinin hedeflerine ve mimarisine kısa bir genel bakış sunulmaktadır.
Hedefler
Günümüzde XLA, çeşitli makine öğrenimi çerçevesi ön ucunu (PyTorch, TensorFlow ve JAX dahil) desteklemektedir ve önde gelen makine öğrenimi donanım ve yazılım kuruluşlarının iş birliğiyle geliştirdiği açık kaynaklı makine öğrenimi derleyici teknolojilerinin oluşturduğu bir ekosistem olan OpenXLA projesinin bir parçasıdır. OpenXLA projesi oluşturulmadan önce XLA, TensorFlow projesinin içinde geliştirilmişti ancak temel hedefler aynı kalıyor:
Yürütme hızını iyileştirin. Kısa ömürlü işlemlerin yürütme süresini kısaltmak ve çalışma zamanı ek yükünü ortadan kaldırmak için alt grafikleri derleyin, bellek ek yükünü azaltmak için ardışık düzen işlemlerini bir araya getirin ve daha agresif sabit yayılım elde etmek amacıyla bilinen tensör şekillerini uzmanlaştırın.
Bellek kullanımını iyileştirin. Birçok ara depolama arabelleğini ortadan kaldırarak bellek kullanımını analiz edin ve planlayın.
Özel işlemlere bağımlılığı azaltın. Otomatik olarak yoğunlaşan düşük seviyeli operasyonların performansını, başlangıçta elle beslenen özel operasyonların performansıyla eşleşecek şekilde iyileştirerek çok sayıda özel işleme olan ihtiyacı ortadan kaldırın.
Taşınabilirliği iyileştirin. Yeni donanımlar için yeni bir arka uç yazmayı nispeten kolay hale getirin. Böylece ML modellerinin büyük bir bölümü, söz konusu donanım üzerinde herhangi bir değişiklik yapmadan çalışabilir. Bu, yeni donanımlar için modellerin bu işlemlerden yararlanmak için yeniden yazılmasını gerektiren bireysel monolitik işlemleri özelleştirme yaklaşımından zıtlıktır.
İşleyiş şekli
XLA derleyicisi, StableHLO'da tanımlanan makine öğrenimi çerçevelerinden model grafiklerini alır ve bunları çeşitli mimariler için makine talimatlarına derler. StableHLO, ML çerçeveleri ile derleyici arasında taşınabilirlik katmanı sağlayan sürümü oluşturulmuş bir işlem grubu (HLO = üst düzey işlemler) tanımlar.
Genel olarak, model grafiğini hedef için optimize edilmiş yürütülebilir bir dosyaya dönüştüren derleme işlemi aşağıdaki adımları içerir:
XLA, StableHLO grafiğinde hedeften bağımsız birçok dahili optimizasyon ve analiz geçişi gerçekleştirir. Örneğin, hesaplama için çalışma zamanı belleği ayırmak amacıyla ÖAM, hedeften bağımsız işlem birleştirme ve arabellek analizi gibi yöntemler kullanılır. XLA, bu optimizasyon aşamasında StableHLO diyalektini de dahili bir HLO diyalektine dönüştürür.
XLA, HLO seviyesindeki optimizasyonlar için HLO hesaplamasını bu kez hedefe özel bilgileri ve ihtiyaçları göz önünde bulundurarak bir arka uca gönderir. Örneğin, GPU arka ucu, özellikle GPU programlama modeli için faydalı olan işlem füzyonları gerçekleştirebilir ve işlemin akışlara nasıl ayrılacağını belirleyebilir. Bu aşamada arka uçlar, belirli işlemleri veya bunların kombinasyonlarını optimize edilmiş kitaplık çağrılarıyla kalıp eşleştirebilir.
Arka uç, hedefe özel kod oluşturma işlemini gerçekleştirir. XLA'ya dahil olan CPU ve GPU arka uçları; düşük düzeyli IR, optimizasyon ve kod oluşturma için LLVM kullanır. Bu arka uçlar, HLO hesaplamasını verimli bir şekilde temsil etmek için gereken LLVM IR'yi yayar ve daha sonra bu LLVM IR'den yerel kod yayması için LLVM'yi çağırır.
Bu işlemde XLA derleyicisi, bazı yeni HW mimarilerini hedeflemek için alternatif bir arka uçta kolayca yuvalanması açısından modülerdir. GPU arka uç, şu anda LLVM NVPTX arka ucu aracılığıyla NVIDIA GPU'ları desteklemektedir. CPU arka ucu, birden fazla CPU ISA'sını destekler.