Bu sayfa, Cloud Translation API ile çevrilmiştir.
Switch to English

XLA için MLIR CodeGen

XLA, HloInstruction üzerinde HloInstruction ve bu gösterim üzerinde birçok optimizasyon gerçekleştirerek hedeflenen cihazlar arasında bunların çoğunu paylaşır. Bir nokta olarak doğrusal bir çizelge hesaplanır ve bellek tamponu her değere statik olarak atanır. Cihaza özel kodgen, bu diziyi hareket ettirerek ve cihaza uygun bir gösterim oluşturmak için "yayıcılar" olarak çağrılarak çalışır (örneğin CPU'daki XLA hesaplaması başına tek bir LLVM işlevi veya GPU işlemlerini ve muhtemelen üretilen PTX dizisini bir dizi hedefleme GPU'su).

Bir aşamalandırma aşaması olarak, şu anda XLA tampon atama aşamasını tamamladıktan ve bunun yerine lhlo lehçesinde bir MLIR modülü yaydıktan hemen sonra işlemi durdurma sürecindeyiz. Buradan, cihaza bağlı olarak MLIR bileşenlerini (esas olarak Linalg, afin ve GPU lehçesi) kullanarak kod oluşturuyoruz.

lhlo girişi olarak lhlo kullanarak XLA / GPU'yu lhlo olarak lhlo için kayıt planı aşağıdadır.

Görevler

evsahibi cihaz
Giriş formatı HloInstruction * (Görev 1) HloInstruction * (Görev 1)
Çıkış formatı xla :: Thunk (Görev 2) LLVM IR (Görev 3)
  • Görev 1 hem ana bilgisayar hem de aygıt giriş biçimini HloInstruction * yerine LHLO olarak değiştirir.
  • Görev 2 ana bilgisayarın çıktı biçimini gövdeden "ana makine için bazı iniş alanı" olarak değiştirir (aşağıya bakın).
  • Görev 3, aygıt çıktısını LLVM IR'den bir MLIR biçimine geçirir. Bu projeye isteğe bağlıdır ve ayrıntılar için "Cihaz Taşıma LLVM IR" bölümüne bakın.

Bu proje, mümkün olduğunca LHLO yayıcıları ile uçtan uca çalıştırılabilir modellere sahip olmaya öncelik vermektedir. Bu, öncelikli olarak aşağıdaki sipariş hedefleri listesinin ima edildiği anlamına gelir:

  • Mevcut Thunks ve yayıcılar değiştirilmeden XLA / GPU'yu LHLO yayıcılarla çalıştırılabilir yapın.
  • Her durumda, LHLO'daki HloInstruction * referanslarını ortadan kaldırın:
    • Eski bir yayıcıyı MLIR tabanlı bir yayıcıya (ör. Linalg) veya
    • MLIR temsilini almak için mevcut vericiyi mekanik olarak çevirin (GPU Dialect ile Standard'a geçin).

Göçmen Gövdeleri (Görev 2)

xla :: gpu :: Thunk bir veri yapısıdır:

  • Ana bilgisayardan çağrılabilir (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Alt sınıflarında çeşitli veriler taşır.
  • BufferAllocation :: Slice ve StreamExecutor ile etkileşime girer.
  • Çekirdekleri başlatır
  • Tüm çalışma zamanı kitaplıklarını çağırır.

Bunun maliyeti şunları içerir:

  • Op'a özel yapılandırma verilerini temsil etme (örneğin, evrişim konfigürasyonları).
  • Op şekli ve işlenen şekilleri geçirme.
  • Bir ağaç gövdesini temsil ederken (süre, koşul vb.)

Geçiş çalışması LHLO / yayıcı göçünden bağımsızdır. Sınırlı kaynaklar altında, LHLO / yayıcı göçünün arkasında önceliklidir.

Ana bilgisayar tarafı parçasının LHLO'dan nasıl düşürüleceğine dair birkaç seçeneğimiz var:

  • TFRT
    • (Pro) kullanım için harika CUDA ve HIP sarmalayıcılar.
    • TFRT ops C ++ kodu ile yorumlandığı gibi (Pro) uygulaması kolay kütüphane çağrıları (cuDNN, cuBLAS, cuFFT, vb).
    • (Con) konak tarafı geliştirilme aşamasındadır ve test edilmemiştir.
  • Jitted CPU kodu
    • (Pro) büyük düşük yetenek. Birkaç döngü ve koşul oluşturun ve bitti.
    • (Con) GPUDialect henüz zincirleri / akışları / eşzamansızlık / cihaz tahsisini modellemiyor.
    • (Con) CUDA / HIP çalışma zamanı desteği minimumdur (araç seti yolu, sürüm, dinamik yükleme, vb.).
  • Mevcut (yorumlama) XLA çalışma zamanı

Karar: TFRT benimsenmesi, aynı zamanda TFRT CPU kodu jitting destekler.

Geçiş Cihazı LLVM IR (Görev 3)

Bir elemental yayıcı, öğeyi tek tek doldurarak hedef op üretir. Her çıkış elemanı, işlenenlerden bir dizi öğeye bağlıdır. Tüm elemanlar, tamponun dinamik indekslerle birleştirilmesiyle açıklanmaktadır. Hemen hemen tüm "matematik" op'larını tanımlamak yeterlidir, ancak performans nedenleriyle doğrudan "matematik" op'lerinin büyük bir alt kümesi doğrudan (Cpu | Gpu) ElementalIrEmitter'da uygulanır.

ElementalIrEmitter benzersizdir:

  • Kodun büyük bir kısmı XLA / GPU ve CPU arasında paylaşılır.
  • Tüm elemanlar arasındaki operasyonlar dahil olmak üzere modellerde görülen op'ların büyük bir bölümünü temsil eder.
  • Çoğu füzyon sadece ElementalIrEmitter'a bağlıdır.
  • Op elemanları ve işlenen elemanları arasındaki veri bağımlılığı DAG'sini açıkladığı için yapısal olarak basittir.
  • Çoğunlukla taşınabilir ve üst düzeydir (örneğin GPU kReduce ve GPU kCopy'den farklı olarak).
  • Dinamik şekil desteği, en azından eleman bazında operasyonlar için kolaydır.

Şimdi, temel olarak yayılan veya yayılmayan tüm oplar için, her XLA opunun son durumunun birkaç çeşidi vardır:

  1. Cihaz kodu LLVM IR olarak kalır.
  2. Eski vericiyi LHLO gibi olacak şekilde yeniden düzenleyin -> MLIR LLVM Lehçesi:
    • (Maliyet) Eğer nihayetinde Standard'a geçmek istersek, işe koyulmayacağız.
    • (Fayda) Kolay ve mekaniktir. Kısa sürede yapılabilir.
    • (Fayda) (1) ile karşılaştırıldığında daha fazla fayda sağlamaz.
  3. Refactor eski yayıcılar LHLO -> MLIR GPU + Standart + Döngüler gibi:
    • (Maliyet) Mevcut yayıcıları Standarda yükseltmek bazı zorluklar doğurur. İşaretçilerin ve GEP'lerin MemRef'lere ve SubView'lere dönüştürülmesi gerekir. Amdgpu tamlığının sağlanması başka bir şeydir.
    • (Maliyet) XLA / GPU, ağırlıklı olarak LLVM meta verilerine dayanır:
      • blok / evre indeksleri için range .
      • align , dereferenceable , invariant.load , alias.scope , noalias yükleme / mağazalar için.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable . sıralı döngüler için llvm.loop.vectorize.enable .
    • (Fayda) Uzun vadeli olabilir. Daha taşınabilir.
  4. Refactor eski yayıcılar LHLO olmak -> Linalg ve yeni Linalg yayıcılar yazmak
    • (Maliyet) Bu duruma göre değişir. Önceki seçeneklerle karşılaştırıldığında, XLA'nın performansıyla eşleşen yeni bir uygulamanın, bazı operasyonlar için önemli bir maliyet olabilen <-> iş akışını optimize etme ölçütünden geçmesi gerekiyor.
    • (Yarar) birleşik yığın; topluluk desteği; taşınabilirlik; daha fazla optimizasyon potansiyeli.

Sonuç:

  • (2) için gitmeyin. (1) veya (3), (2) 'den daha iyidir. (2) maliyeti (1) 'den daha fazladır, çünkü çok fazla mekanik yeniden düzenleme gerektirir. (1) ile hala XLA'nın MLIR yayıcılarını almasını sağlama hedefine ulaşabiliriz. Bu LHLO -> LLVM IR -> eski cihaz yayıcılarını çalıştırmaktır.
  • ElementalIrEmitter ops (4) 'ü seçer, ancak artımsal olarak olmaz. Op ile op yapmanın bir yolu yoktur, çünkü tüm temel olarak yayılan op'lar aynı grafiğe bağlanır. Bu çalışma aynı zamanda devam eden birkaç kuvvetin (xla / service / mlir_gpu, çekirdek üreteci Linalg) birleşme noktası olarak da kullanılabilir.
  • Diğer tüm operasyonlar (1) için geçerlidir. Esnek bir hedef olarak, (3) veya (4) 'e geçebilirler.

Önceliklendirme

Yukarıda belirtilen üç görev de paralelleştirilebilirken, sınırlı kaynaklar altında serileştirilmeleri gerekir. Önceliklendirme her bir görevin tamamlanması için görünür sonuçlara odaklanır.

Öncelik şu şekildedir: Görev1 (eski yayıcılar için LHLO)> Görev 2 (Thunks)> Görev 3 (MLIR yayıcılar).

Görev 1'in sonunda, XLA kullanıcıları bir LHLO (örneğin çekirdek üreteci) üretebilir ve bunları yürütebilir. Derleme formatı serileştirilemez MLIR.

Görev 2'nin sonunda, LHLO uygun, serileştirilebilir MLIR'a düşer. Bu, çevrimdışı derlemeyi etkinleştirir.

Görev 3'ün sonunda, tüm XLA yayıcılar uygulamada MLIR tabanlıdır.

Detaylı tasarım

Adım 1: (Görev 1) LHLO'yu tamamlayın ve Eski Yayıcıların LHLO Almasını Sağlayın

Bu adım, mevcut tüm XLA / GPU yayıcıların MLIR ops ile etkileşime girmesini sağlar. Bu adım saf yeniden düzenleme ve NFC'dir.

Bu adım çoğunlukla mekaniktir, ancak dürüst olmayan bir HloComputation ve LHLO arasındaki aşağıdaki tutarsızlıkları fark etmeye değer:

  • Her HloInstruction kendi işlenenlerine (veri akışı DAG) doğrudan erişime sahiptir. Aksine, her LHLO op sadece operand tamponlarına (ops ve tamponlar arasında iki taraflı) erişebilir. LHLO operasyonları, işlenen operasyonlarına erişmek için kullanım-def zincirlerinden geçmelidir.
  • Unnested eski miraslar ampirik olarak neredeyse hiçbir zaman işlenenlerine erişmezler. Tek istisna kReduce.
  • Sınırsız eski yayıcılar bufferAssignment'a yalnızca dilimleri almak için erişir, dataflow_analysis () veya alias_analysis () gibi aux veri yapılarına erişmek için değil. llvm_ir dilim bilgilerine dayanarak kendi takma_analizini () oluşturur.

Sonuç, LHLO'nun büyük bir güçlük çekmeden sağa oturması gerektiğidir.

2. Adım: (İsteğe bağlı) Profil Oluşturma Desteği

Bu adım yalnızca bazı XLA Thunk mantığını atmaya başlarsak gereklidir (bir sonraki adıma bakın).

Herhangi bir MLIR tabanlı yayıcı açmadan önce, MLIR tabanlı yayıcılar için profil oluşturmamız gerekir.

Şu anda XLA, StreamExecutor'ın zamanlayıcısını arayarak kendi profilini oluşturuyor. Kaputun altındaki zamanlayıcı, çekirdek başlatmadan önce ve sonra iki olay ekler ve bu iki olay arasındaki senkronizasyon süresini ölçer.

MLIR'de profil oluşturmayı desteklemek için kabaca üç yaklaşım vardır:

  • Profil oluşturucuyu uçtan uca çalıştırın
  • Enjekte edilmiş profil oluşturucu kullanarak LHLO'daki her bir işlem için bir profil op ekleyin.

"Uçtan uca" yaklaşım MLIR için şeffaftır, ancak XLA'nın bunu ilk etapta kullanmamasına neden olan aynı sorunla karşı karşıyadır: bir profiler (nvprof / ...) tarafından toplanan kütüphane çağrıları HLO ile kolayca ilişkilendirilemez ops. Örneğin, cuDNN her HLO için birden çok çekirdek başlatır ve hangi çekirdeklerin hangi HLO'ya karşılık geldiğini söylemek zordur.

"Enjekte edilen profil oluşturucu" yaklaşımı şunları gerektirir:

  • LHLO bir profil oluşturucuyu parametre olarak alır.
  • profile ekleyerek start / profile.end her op önce ve sonra.
  • bu profilden bir geçiş. {start, end} bir C ++ uygulamasına geçer.

Tam profil oluşturma, MLIR tarafından oluşturulan operasyonlar için kolayca yapılamaz, çünkü:

  • MLIR'in bir zamanlayıcısı yoktur ve TFRT / StreamExecutor'a da bağlıdır.
  • MLIR, karmaşık parametrelerle C fonksiyonlarına kolayca çağrı yapmaz.

Adım 3: (Görev 2) Gömme Thunks

Not olarak, kabaca üç tür thunk vardır:

  • Çekirdek: Bir çekirdeği başlatan thunk.
  • Ana bilgisayar kontrol akış mantığına sahip (kontrol dizisi için koşullu, sıralı) ve gövde çekirdeklerini başlatan kontrol akış gövdeleri.
  • Kütüphane başlıkları: cuDNN, cuBLAS, cuFFT, NCCL, vb.

Plan:

  • Thunks (de) serileştirilebilir yapın.
  • TFRT'nin bu anlambilimi destekleyebileceği bir duruma gelmesine yardımcı olun.
  • Devlet geliştikçe bireysel gövdeleri kademeli olarak geçirin.

Bu işlem öğeleri yalnızca kısmen sıralanmıştır. Gerçek yürütme emri / mühendislik paralellikleri olduğu gibi değerlendirilecektir.

Adım 4: (Görev 3) Taşınan ElementalIrEmitter

Profil oluşturma hazır olduğunda, MLIR'deki tüm ElementalIrEmitter tabanlı yayıcıları tamamlayabilir ve ayarlayabiliriz. Ardından, bu MLIR tabanlı yayıcıların hepsinin tek bir akış kullandığını varsayarak varsayılan olarak açarız.

Kodun büyük bir bölümünü paylaştıkları için XLA / CPU'nun ElementalIrEmitter'ı da taşımanın yararlı olduğuna dikkat edin.

Tüm kıyaslama ve performans araştırmaları yapıldıktan sonra (TODO: performans paritesini tanımlayın), yeni MLIR tabanlı temel vericiyi açıyoruz ve eski ElementalIrEmitter'ı siliyoruz.

Bu adım ayrıca daha sonraki geçiş için kolay füzyon geçişleri (iç içe ops) sağlar.

Adım 5: Çoklu Akış Desteği veya Bırakma

MLIR'da desteklemedikçe veya özelliği bırakana kadar yayıcıların bazılarını silemeyiz. MLIR'da nispeten büyük miktarda çalışma ve XLA için az miktarda kazanç. Çok akışlı XLA / GPU kullanıcılarının mevcut kullanıcılarını araştırmalı ve makul olması durumunda bu özelliği silmeye çalışmalıyız.

Adım 6: (Görev 3) Taşınan Cihaz İşlemleri

Bu adım, tüm dürüst olmayan operasyonları geçirir, ardından tüm dürüst olmayan yayıcıları silebiliriz.

Bu, kCopy ve kReduce için bir yeniden yazma / yeniden düzenleme çağırır. kReduce zaten bolca üzerinde çalışılıyor, bu yüzden yapılması gereken gerçek iş miktarı görülmeye devam ediyor.