ML Topluluk Günü 9 Kasım! TensorFlow, JAX güncellemeler için bize katılın ve daha fazla bilgi edinin

XLA Mimarisi

XLA'yı neden yaptık?

XLA'nın TensorFlow ile çalışması için birkaç hedefimiz vardı:

  • Yürütme hızını artırın. TensorFlow çalışma zamanından kaynaklanan ek yükü ortadan kaldırmak için kısa ömürlü Ops yürütme süresini azaltmak için alt grafikleri derleyin, bellek ek yükünü azaltmak için ardışık düzenlenmiş işlemleri birleştirin ve daha agresif sabit yayılmaya izin vermek için bilinen tensör şekillerinde uzmanlaşın.

  • Bellek kullanımını geliştirin. Prensipte birçok ara depolama arabelleğini ortadan kaldırarak bellek kullanımını analiz edin ve planlayın.

  • Özel Operasyonlara olan bağımlılığı azaltın. Otomatik olarak kaynaştırılmış düşük seviyeli Operasyonların performansını, elle birleştirilen özel Operasyonların performansıyla eşleşecek şekilde iyileştirerek birçok özel Operasyon ihtiyacını ortadan kaldırın.

  • Mobil ayak izini azaltın. Önceden alt grafiği derleyerek ve doğrudan başka bir uygulamaya bağlanabilen bir nesne / başlık dosyası çifti göndererek TensorFlow çalışma zamanını ortadan kaldırın. Sonuçlar, mobil çıkarım için ayak izini birkaç kat azaltabilir.

  • Taşınabilirliği artırın. Yeni donanım için yeni bir arka uç yazmayı nispeten kolaylaştırın, bu noktada TensorFlow programlarının büyük bir kısmı bu donanım üzerinde değiştirilmeden çalışacaktır. Bu, TensorFlow programlarının bu Ops'lardan yararlanmak için yeniden yazılmasını gerektiren yeni donanım için bireysel monolitik Ops'u uzmanlaştırma yaklaşımıyla tezat teşkil ediyor.

XLA nasıl çalışır?

XLA'nın giriş diline "HLO IR" veya sadece HLO (Yüksek Seviye İşlemler) denir. HLO'nun semantiği, Operation Semantics sayfasında açıklanmıştır. HLO'yu bir derleyici IR olarak düşünmek en uygun yoldur.

XLA, HLO'da tanımlanan grafikleri ("hesaplamalar") alır ve bunları çeşitli mimariler için makine talimatlarında derler. XLA, bazı yeni HW mimarisini hedeflemek için alternatif bir arka uca yerleştirmenin kolay olması açısından modülerdir. X64 ve ARM64 için CPU arka ucu ve NVIDIA GPU arka ucu, TensorFlow kaynak ağacında yer alır.

Aşağıdaki şema, XLA'daki derleme sürecini gösterir:

XLA, CSE , hedeften bağımsız işlem füzyonu ve hesaplama için çalışma zamanı belleği ayırmak için arabellek analizi gibi hedeften bağımsız birkaç optimizasyon ve analiz geçişiyle birlikte gelir.

Hedeften bağımsız adımdan sonra XLA, HLO hesaplamasını bir arka uca gönderir. Arka uç, bu sefer hedefe özel bilgiler ve ihtiyaçlar göz önünde bulundurularak HLO düzeyinde daha fazla optimizasyon gerçekleştirebilir. Örneğin, XLA GPU arka ucu, özellikle GPU programlama modeli için faydalı işlem füzyonu gerçekleştirebilir ve hesaplamanın akışlara nasıl bölüneceğini belirleyebilir. Bu aşamada, arka uçlar ayrıca belirli işlemleri veya bunların kombinasyonlarını optimize edilmiş kütüphane çağrılarıyla model eşleştirebilir.

Bir sonraki adım, hedefe özel kod üretmektir. XLA'ya dahil olan CPU ve GPU arka uçları, düşük seviyeli IR, optimizasyon ve kod üretimi için LLVM kullanır. Bu arka uçlar, XLA HLO hesaplamasını verimli bir şekilde temsil etmek için gerekli LLVM IR'yi yayar ve sonra bu LLVM IR'den yerel kodu yaymak için LLVM'yi çağırır.

GPU arka ucu şu anda LLVM NVPTX arka ucu aracılığıyla NVIDIA GPU'ları desteklemektedir; CPU arka ucu birden çok CPU ISA'yı destekler.