XLA mimarisi

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

Bugün XLA, çeşitli ML çerçevesi ön ucunu (PyTorch, TensorFlow ve JAX dahil) desteklemektedir ve lider makine öğrenimi donanım ve yazılım kuruluşlarının iş birliğiyle geliştirdiği açık kaynaklı makine öğrenimi derleyici teknolojilerinden oluşan 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ır:

  • Yürütme hızını artırın. Kısa ömürlü işlemlerin yürütme süresini kısaltmak için alt grafikleri derleyin, çalışma zamanından kaynaklanan ek yükü ortadan kaldırın, bellek ek yükünü azaltmak için ardışık düzen işlemlerini birleştirin ve daha agresif sabit yayılım sağlamak amacıyla bilinen tensör şekillerini uzmanlaştırın.

  • Bellek kullanımını iyileştirin. Birçok ara depolama arabelleğine gerek kalmadan bellek kullanımını analiz edip planlayın.

  • Özel işlemlere bağımlılığı azaltın. Otomatik olarak birleştirilen düşük seviyeli operasyonların performansını, başlangıçta elle kaynaştırılan özel operasyonların performansıyla eşleşecek şekilde iyileştirerek birçok ö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 daha kolay hale getirin. Böylece, ML modellerinin büyük bir bölümü bu donanım üzerinde değiştirilmeden çalışabilir. Bu yaklaşım, yeni donanımlar için bağımsız monolitik işlemleri uzmanlaşma yaklaşımından farklıdır. Yeni donanımlar için modellerin bu işlemlerden yararlanması için yeniden yazılması gerekir.

İş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ında 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 yürütülebilir dosyaya dönüştüren derleme işlemi aşağıdaki adımları içerir:

  1. XLA, Kararlı HLO grafiğinde işlem için çalışma zamanı belleğini ayırmak üzere ÖAM, hedeften bağımsız işlem füzyonu ve arabellek analizi gibi hedeften bağımsız birçok yerleşik optimizasyon ve analiz geçişi gerçekleştirir. Bu optimizasyon aşamasında XLA, StableHLO diyalektini de dahili bir HLO diyalektine dönüştürür.

  2. XLA, HLO düzeyinde daha fazla optimizasyon için HLO hesaplamasını bir arka uca gönderir. Bu kez, hedefe özel bilgiler ve ihtiyaçlar göz önünde bulundurulur. Örneğin, GPU arka ucu, özellikle GPU programlama modeli için faydalı olan işlem füzyonları gerçekleştirebilir ve hesaplamanın 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.

  3. 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 sanal makinesini 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 yaymak 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 yer alabilmesi açısından modülerdir. GPU arka ucu, şu anda LLVM NVPTX arka ucu aracılığıyla NVIDIA GPU'ları desteklemektedir. CPU arka ucu birden fazla CPU ISA'sını destekler.