Bu belgede, bir XLA Üst Düzey Optimizasyon (HLO) modülünün başlangıç durumundan son yürütülebilir duruma kadar olan yolculuğu özetlenmektedir. Bazen "modül"ü çıkarıp yalnızca "HLO" olarak bahsederiz.
Ön optimizasyon HLO
Ön optimizasyonlu HLO modülüyle başlıyoruz. Ön optimizasyon HLO'su, XLA'ya özgü olduğu düşünülen işlemler (ops) içermez. Örneğin, fusion veya bitcast. Bu aşamada işlemlerin düzeni yoktur veya varsa bu düzen yoksayılır. Ön optimizasyon HLO'su genellikle TensorFlow ve JAX gibi üst düzey çerçeveler tarafından üretilir. XLA işareti -xla_dump_to kullanılırken ön optimizasyon HLO'su, "before_optimizations.txt" dosya adı sonekiyle bir dosyaya aktarılır.
HLO modülünü optimize etme
XLA:GPU ardışık düzeni, bir dizi geçiş çalıştırarak ön optimizasyon HLO'sunu optimize edilmiş HLO'ya dönüştürür. Kartlar semantik olarak gruplandırılabilir ve aşağıdaki sırayla çalıştırılabilir:
Parçalama ile ilgili kartlar
Shardy Partitioner gibi geçişler veya SPMD parçalama için olanlar bu kapsamdadır.
Optimizasyon geçişleri
Bu kapsamda hem yasallaştırma hem de basitleştirme işlemleri yer alabilir.
Toplu optimizasyon geçişleri
Optimizasyon geçişlerine benzer ancak toplu işlemlere odaklanır.
Düzen atama kartları
Her HLO işlemine, talimat şeklinin bir parçası olan bir düzen atanır. Düzen, tensörün bellekte fiziksel olarak nasıl düzenleneceğini kontrol eder.
Düzen içeren şekil örneği:
f32[10,20,30]{2,0,1}
Öğe türünden sonra şeklin mantıksal boyutları ve ardından küçükten büyüğe doğru düzen permütasyonu gelir. Bu örnekte en küçük boyut 30, ikinci en küçük boyut 10 ve büyük boyut 20'dir.
Düzen atamanın amacı, açgözlü bir strateji kullanarak gerekli fiziksel transpozisyon sayısını en aza indirmektir. Belirli düzen kısıtlamalarıyla (ör. cuDNN/cuBLAS kitaplıkları, ardışık boyutlar bekler) başlar ve düzenleri HLO grafiğinde önce "aşağı" sonra "yukarı" doğru yayar. Düzen yayılımının sonunda, bazı talimatların çakışan düzenleri olabilir. Bunlardan biri bir işlenenden, diğeri ise bir kullanıcıdan yayılır. Bu çakışmayı çözmek için, düzeni işlenen düzeninden talimat düzenine değiştiren bir copy HLO talimatı eklenir.
Düzen normalleştirme geçişleri
Fiziksel şekli anlamak biraz zor olduğundan düzen normalleştirme, şekli varsayılan düzeni {rank-1, rank-2, …, 0} kullanacak şekilde yeniden yazmaya çalışır. Yukarıdaki örnekte, normalleştirilmiş şekil f32[20,10,30]{2,1,0} olur. Düzenleri değiştiren kopyalama işlemleri, transpose ve bitcast kombinasyonu olarak yeniden yazılır. Şu anda tüm işlemleri normalleştiremediğimiz için varsayılan olmayan düzenlere sahip olabilecek bazı işlemler vardır. Bunların en önemlileri gather ve dot işlemleridir. Normalleştirilmiş işlemler ile normalleştirilmemiş işlemler arasındaki sınırlarda, transpozu temsil eden bitcast işlemler olacaktır. Yani, fiziksel olarak işlem yapmayan bir düzen atanmış bir transpoz.
Düzen normalleştirme, bazı örtülü transpozeleri de açık hale getirir. Bu, kod oluşturma işlemi açık transpozeleri özel bir yayıcıyla işleyebildiğinden önemlidir. Örneğin, yeniden şekillendirme işleminin, işlenen ve sonuç arasında farklı bir fiziksel düzene sahip olmasına (ör. farklı sıralama nedeniyle) teknik olarak izin verilir. Düzen normalleştirme geçişlerinin bir parçası olarak çalışan ReshapeDecomposer geçişi, yeniden şekillendirmeyi transpose, yeniden şekillendirme bitcast ve transpose dizisine dönüştürür.
Düzen atama optimizasyonu geçişlerini yayınlama
Buradaki en önemli geçişler Triton birleştirme işlemleri (GEMM birleştirme işlemleri + Softmax/Layernorm birleştirme işlemleri) veya kitaplık çağrılarına yeniden yazma işlemleridir. Otomatik ayarlama bu adımda da çalışır. Burada XLA, farklı yayıcılar arasında seçim yapar, kıvrımlar veya noktalar için en iyi algoritmayı seçer, Triton yayıcı tarafından işlenen birleştirmeler için en iyi döşemeyi bulur.
Karma kartlar
İki ana geçiş PriorityFusion ve Multi-Output füzyonudur.
PriorityFusion bölümünde, maliyet modeline göre birleştirme işlemleri yaparız. Birleştirme sırasında, işlem tüm kullanıcılarla birleştirilebiliyorsa birden fazla kullanıcıyla işlemleri çoğaltmaya izin verilir. Mümkünse mevcut Triton Softmax füzyonlarının genişletilmesine de izin verilir.
Multi-Output fusion, bir işleneni paylaşan işlemleri/birleştirme işlemlerini birleştirmeye olanak tanıyan ayrı bir geçiştir. Ayrıca, ek çıkışlar ekleyerek işlemleri/işlem birleştirme işlemlerini kullanıcılarla birleştirebilir. Böylece, birleştirilecek işlemin diğer kullanıcıları bu çıkışlara yönlendirilebilir. Bu geçişin, HLO grafiğine döngü eklememeye dikkat etmesi gerekir.
Çoklu çıkış birleştirme işleminden sonra, ortak alt ifadeyi ortadan kaldırma (HloCSE geçişi) çalıştırılır. Bu işlem, daha önce yinelenen işlemleri aynı birleştirme işleminde yer aldılarsa tekrar birleştirebilir.
Birkaç birleştirme sonrası geçiş
Kolektiflerle ilgili çeşitli geçişler (ör. bunları eşzamansız hale getirme veya kolektiflerin belirli bir göreceli sırasını zorunlu kılma).
Son olarak, yerinde işlemlerin başka yerlerde hâlâ ihtiyaç duyulan verilerin üzerine yazılmamasını sağlamak için kopyaların eklendiği CopyInsertion işlemini çalıştırırız.
Optimizasyonun sonunda, -xla_dump_to işareti kullanılıyorsa optimize edilmiş HLO, dosya adı soneki "after_optimizations.txt" olan bir dosyaya aktarılır. HLO'yu HloModule'ü gerçekten değiştiren ara geçişlerden sonra boşaltmak istiyorsanız -xla_dump_hlo_pass_re=.* işaretini (veya belirli geçişlerle sınırlamak için belirli bir normal ifade) kullanabilirsiniz.
Planlama
Programı olmayan bir HLO modülü, işlemlerin işlenme sırası konusunda yine de bir miktar serbestliğe sahiptir. İşlenen/sonuç ilişkilerine ve kontrol bağımlılıklarına uyan herhangi bir topolojik sıralama geçerlidir. Planlama, hangi belirli sıranın kullanılacağını belirler. Bu aşamadaki temel endişe, tensörlerin kullanım ömrüne bağlı olan maksimum bellek tüketimidir. İlk adımda, farklı planlayıcı algoritmalarını deneriz ve en yüksek bellek tüketimini en aza indirecek planı seçeriz. Bu noktada henüz fiziksel arabelleklerle çalışmadığımızı (bu, "Arabellek Atama" bölümünde gerçekleşir) ve bellek kullanımını simüle ettiğimizi unutmayın.
Ardından LatencyHidingScheduler geçişi çalışır ve hesaplama-iletişim çakışmasını en üst düzeye çıkarmaya çalışır. Ancak bu durum, bellek kullanımını tekrar artırabilir.
Son olarak, en yüksek bellek tüketimi, elimizdeki bellek miktarından yüksekse HloRematerialization komutunu çalıştırırız. Bu geçiş, performansı düşürerek bellek kullanımını azaltmaya çalışır. Örneğin, bazı birleştirmeler bölünebilir ve bazı işlemler, daha kısa arabellek ömrüne sahip olmak için çoğaltılabilir. Yeniden materyalleştirme gerçekleşirse model tarafında bellek gereksinimlerini azaltmanın yollarını (ör. daha küçük toplu iş boyutları kullanma) araştırmanız faydalı olabilir.
Arabellek Ataması
LLVM IR'ye düşürmeden hemen önce, HLO grafiğindeki her talimata arabellek dilimleri atayan arabellek atama geçişlerini çalıştırırız. Arabellek atama işlemi birkaç adımda gerçekleştirilir:
Talimatlara
HloDataflowAnalysisatarHloValues(temelde mantıksal arabellekler). Yerinde işlemler için bir işleneninHloValueyeniden kullanılabilir. Bir işlem, birden fazlaHloValuetanımlayabilir (ör. demet sonuç şekliyle).HloAliasAnalysis, diğer ad işlemleri için arabellekleri birleştirmeye çalışır veHloValueileHloBufferarasında bir eşleme hesaplar.BufferAssignment,HloBuffersöğesinin büyük bir arabellekteki arabellek dilimleriyle eşlemesini, aynı arabellek diliminin, yaşam süreleri çakışan farklıHloBuffersöğeleri için kullanılmayacağı şekilde hesaplar. Diğer ad oluşturabilecek işlemler için küçük bir çakışma olması (birHloBufferişleminin bitiş zamanı, diğerHloBufferişleminin başlangıç zamanıyla çakışabilir) sorun değildir.-xla_dump_toişareti kullanılırken arabellek atamasıyla ilgili bazı bilgiler, "after_optimizations-buffer-assignment.txt" adlı dosyanın sonuna eklenir.
Thunks
Bir HLO grafiği optimize edilip planlandıktan sonra belirli bir arka uç (CPU veya GPU) için thunk'ların doğrusal bir dizisine indirgenir.
XLA'da Thunk, çalışma zamanının yürüttüğü bağımsız bir iş biriminin soyutlamasıdır. Derlenmiş bir çekirdek başlatma, belirli bir işlem, kitaplık çağrısı, kontrol akışı yapısı, toplu iletişim vb. olabilir. Thunk Sırası, belirli bir arka uç için yürütülebilir dosyanın tamamını temsil eder.
Thunk Emission
Planlanmış bir HLO hesaplamasını thunk dizisine dönüştürme işlemine "thunk yayımı" adı verilir. Bu, her arka uçtaki özel bir yayıcı sınıfı tarafından işlenir.
GPU arka ucunda bu işlem IrEmitterUnnested tarafından gerçekleştirilir.
EmitHloComputation, bir hesaplamada planlanmış HLO talimatları listesini yineler ve özel bir Emit... yöntemine (ör.
EmitFusion, EmitConvolutionThunk, EmitWhile). Bu yöntemlerin her biri uygun Thunk nesnelerini oluşturur ve bunları thunk dizisine ekler.
CPU Backend için ThunkEmitter bu rolü üstlenir ve benzer şekilde düzenlenir. Son ThunkSequence, CpuExecutable içine yerleştirilir.
Bir HLO modülünün giriş hesaplamasındaki her talimatın, son thunk dizisinde sıfır (kTuple, kConstant, ..), bir veya birden fazla (örneğin, sıralama talimatı) thunk'a karşılık gelebileceğini unutmayın.
Komut arabellekleri: GPU'da yürütmeyi optimize etme
Modern GPU donanımı, bir GPU işlemler dizisinin (çekirdek başlatma, bellek kopyalama vb.) bir kez kaydedilmesine ve ardından dizinin minimum CPU yüküyle birden çok kez oynatılmasına olanak tanır. Bu, özellikle çok sayıda küçük ve hızlı başlatılan çekirdeğe sahip iş yükleri için kritik bir performans optimizasyonudur. XLA, CUDA grafikleri veya HIP grafiklerinin soyutlaması olarak Komut Arabelleği'ni kullanır. Temel arayüz, GpuCommandBuffer içinde tanımlanır.
Komut arabelleği, CommandBufferThunk ile bir thunk dizisinde gösterilir.
Yayıcı, bu thunk'ı doğrudan HLO talimatlarından oluşturmaz. Bunun yerine, bu işlem ThunkSequence üzerinde çalışan CommandBufferConversionPass tarafından yapılır.
Bu geçiş, uyumlu thunk'ların bitişik alt dizilerini (ör. bir dizi KernelThunk ve GemmThunk) tanımlar. Ardından, bulunan alt diziyi tek bir CommandBufferThunk ile değiştirir. Yeni thunk, orijinal thunk'ların mantığını hafif CommandBufferCmd nesnelerinin listesi olarak kapsar.
Belirli bir GPU akışında ilk kez yürütülen bir CommandBufferThunk, komut dizisini bir donanım komut arabelleğine "kaydeder". Sonraki tüm yürütmelerde, kaydedilen diziyi "tekrar oynatmak" için GPU'ya tek bir komut gönderir. Bu sayede, her bir çekirdeğin ayrı ayrı başlatılmasıyla ilgili CPU ek yükü önlenir.
Yürütülebilir
XLA derleme işlem hattının nihai ürünü, bağımsız ve platforma özel bir çalıştırılabilir dosyadır. Bu nesne, derlenmiş programın hedef cihazda (ör. CPU veya GPU) çalıştırılması için gereken tüm bilgileri kapsar. Derleyici ile çalışma zamanı arasındaki köprüdür. PJRT gibi modern çalışma zamanları biraz daha yüksek düzeyde soyutlamalar kullanır (bkz. PjRtExecutable). Ancak bunlar sonuçta arka uca özgü bir yürütülebilir dosyayı sarmalar.
Executable, derleme sırasında oluşturulan çeşitli önemli bilgiler içerir. Tam içerik arka uca göre değişse de genellikle şunları içerir:
Derlenmiş Kod: Bu, cihazda çalışacak düşük seviyeli makine kodudur. CPU'lar için bu genellikle bir veya daha fazla nesne dosyasıdır. GPU'lar için bu, çalışma zamanında GPU'ya yüklenen PTX veya HSACO biçimindeki derlenmiş cihaz kodudur.
Yürütme planı (ThunkSequence): Çalışma zamanı mantığının çekirdeği. Bu, Thunk nesnelerinin doğrusal bir dizisidir. Her thunk, tek bir iş birimini (ör.çekirdek başlatma, kitaplık işlevi çağırma (ör. cuBLAS) veya kontrol akışını işleme) temsil eder. Çalışma zamanı, bu sırayı yineleyerek programı yürütür.
Bellek Düzeni (BufferAssignment): BufferAssigner tarafından üretilen bu kritik meta veri parçası, hesaplamanın tam bellek düzenini açıklar. Her arabelleğin boyutunu ve parametreler, çıkışlar ve geçici değerler için belleğin nasıl ayrıldığını ve yeniden kullanıldığını belirtir. Çalışma zamanı, cihaz belleğini ayırmak ve her bir thunk'a doğru işaretçileri iletmek için bunu kullanır.
(isteğe bağlı) HLO Modülü: Hata ayıklama ve profil oluşturma için yürütülebilir dosya genellikle derlendiği son, optimize edilmiş HloModule'e bir referans tutar.
Nihai yürütülebilir dosyanın oluşturulması, her bir belirli arka uç için derleyici tarafından düzenlenir. Derleyici uygulamasının RunBackend yöntemi, derleme sürecindeki son adımdır. Bu adımda, derlenen tüm yapılar bir Executable nesnesinde paketlenir.
GpuCompiler ve CpuCompiler sırasıyla GPU ve CPU'yu hedefler.
Bir kullanıcı, yürütülebilir bir dosyada Execute... işlevini çağırdığında çalışma zamanı, bellek ayırmak için BufferAssignment işlevini kullanır ve ardından derlenmiş kodu kullanarak cihazda işlemleri başlatmak için ThunkSequence işlevini çağırır.