Gelişmiş hesaplama teknolojileri, yapay zeka ve bilimsel hesaplamalar alanındaki uygulamaların performansını artırmaya devam ederken, GPU optimizasyon tekniklerine olan ilgi yeniden artmıştır. Bir uygulama geliştiricisi olarak, GPU’lar üzerinde program yazmanın birçok yolu bulunmaktadır. Bu yazıda, yazılım yığınını oluşturan farklı seviyeleri tanıtıyor ve bunların en alt seviyesine, yani Parallel Thread Execution (PTX) kodunu yazmaya dalıyoruz.

Hızlandırılmış Hesaplama Yazılım Yığı

Bugün, GPU üzerinde GPU’ya özgü kod yazmadan birçok şey yapabilirsiniz. Kütüphane geliştiricileri ve yazılım mühendisleri, bu alanda önemli işler yapmışlardır. Örneğin, mavi baskılar kullanarak tam yapay zeka iş akışları oluşturabilir veya PyTorch gibi çerçeveler kullanarak modelinizi tanımlayıp, uygun GPU kodu ve kütüphanelerin otomatik olarak kullanılmasıyla programınızı çalıştırabilirsiniz.

Ayrıca, kuantum bilgisayarlama, veri işleme, fizik AI, gen dizileme, kenar hesaplama, ilaç keşfi gibi alanlara özgü kütüphaneleri içeren kapsamlı NVIDIA CUDA-X kütüphanelerini kullanarak uygulamanızı geliştirebilirsiniz. Eğer ihtiyaç duyduğunuz tüm işlevsellik bu alanlara özgü kütüphanelerde yoksa, GPU’ları OpenACC gibi derleyici direktifleri ve libcu++ kullanarak C++ standart kütüphaneleri ile programlayabilirsiniz.

Bütün bu örneklerde, GPU’ya özgü bir kod yazmıyorsunuz; bunun yerine, uzman mühendisler tarafından titizlikle tasarlanmış, uygulanmış ve optimize edilmiş kütüphanelere veya derleyici direktiflerine güveniyorsunuz.

Ancak, bir kütüphane ihtiyaç duyduğunuz işlevselliği sunmuyorsa, o zaman yazılım yığınının daha derinlerine inip yüksek düzey dillerde, örneğin C++, Fortran veya Python gibi, doğrudan CUDA GPU kodu yazmanız gerekebilir.

Son olarak, nadir durumlarda geliştiriciler, performans açısından hassas kod bölümlerinin son derece yüksek performans gerektirdiğini düşünerek doğrudan PTX yazmayı tercih edebilirler. Performans optimizasyon teknikleri ile ilgili olarak, ne kadar kontrol istiyorsanız, yığının o kadar altına inmeniz gerektiğini unutmayın. Bu takasa dikkat edilmelidir: Elle yazılmış düşük seviye kodun sağladığı performans kazançlarının diğer GPU mimarilerine taşınamama durumu, ek geliştirme ve hata ayıklama karmaşıklığını beraberinde getirebilir.

Önceki bir yazımızda gösterdiğimiz gibi, PTX, GPU’ların montaj dilidir. Doğrudan PTX yazmak, çoğu geliştirici için gerekli olmayan son derece ileri düzey bir optimizasyon tekniğidir ve son çare olarak düşünülmelidir. Bununla birlikte, PTX yazmanın sağladığı ince ayar kontrolü, belirli uygulamalarda performans iyileştirmelerine olanak tanır. Bu durumlar genellikle uygulamanın çok performans hassas bölümlerinde, her bir performans artışının önemli avantajlar sağladığı senaryolardır. Mevcut tüm PTX talimatlarını PTX ISA belgesinde bulabilirsiniz.

Bu blog yazısında, el ile yazılmış PTX’in belirli bir AI model uygulamalarında önemli bir algoritmanın performansını nasıl artırdığına dair bir örneği daha derinlemesine inceleyeceğiz.

PTX Yazma Yöntemleri

Örnekleme geçmeden önce, uygulamanıza el yazısı PTX kodunu dahil etmenin bazı yollarını sıralayalım. Yani, ilke olarak bunu nasıl yapabileceğinizi gözden geçirelim. Aşağıda gösterilecek örnek gerçek bir durumu yansıtır ve performans değişikliklerini gösterir.

Satır İçi PTX Kullanımı

PTX kodunu uygulamanıza dahil etmenin standart bir yolu, satır içi PTX kullanmaktır. Bu, aşağıda göstereceğimiz yöntemdir ve sözdizimi ile anlamına dair detaylı bilgiler belgelendirmede mevcuttur. Bu, bir CPU üzerinde montaj kodu yazmaya oldukça benzer.

cuda::ptx Ad Alanı

PTX’i kodunuza dahil etmenin başka bir yolu, libcu++ kullanmaktır; bu kütüphane cuda::ptx ad alanını içerir ve PTX talimatlarına doğrudan karşılık gelen işlevler sunar. Bu, belirli PTX talimatlarını C++ uygulamanız içinde kolayca kullanmanıza olanak tanır. cuda::ptx ad alanı hakkında daha fazla bilgi için belgelere göz atabilirsiniz.

CUTLASS Örneği

Eldeki örneği göstermek adına, lineer cebir alanından belirli bir örnek vereceğiz. Genel olarak, işleminiz bir GEMM olarak ifade edilebiliyorsa, NVIDIA CUBLAS önerilen yöntemdir. CUBLAS, birçok matris boyutu ve şekli için son derece optimize edilmiştir ve birden fazla sayısal hassasiyet seçeneği sunar.

Bazı durumlarda yapmak istediğiniz, CUBLAS’taki işlevsellik ile tam olarak ifade edilmez veya bir GEMM işleminden önce veya sonra doğrudan hesaplamalar yapmanız gerebilir. Bazen, bazı işlevleri çağırmaktan, sonra CUBLAS’tan ve ardından yine bazı işlevleri çağırmaktan ziyade, başka işlemleri bir GEMM işlemi ile birleştirerek performansı artırabilirsiniz. Bu durumun birçok avantajı vardır, çünkü birleşik çekirdekler, örneğin verileri daha etkili kullanmak gibi daha fazla optimizasyon yapmayı mümkün kılabilir.

Bu noktada NVIDIA CUTLASS kütüphanesi devreye girmektedir. CUTLASS, CUDA C++ şablon soyutlamaları içeren, yüksek performanslı matris-matris çarpımını (GEMM) ve ilgili hesaplamaları CUDA’nın her seviyesinde uygulamak için bir koleksiyondur. CUTLASS, GEMM ve benzeri işlemler etrafında daha fazla kontrol ve özelleştirme imkanı sağladığı için geliştiriciden, CUBLAS’tan daha fazla kod yazmasını gerektirir.

CUTLASS, her GPU mimarisi üzerinde mümkün olan en iyi performansı hedefleyecek şekilde tasarlanmış oldukça fazla el yazısı PTX içermektedir. Bu nedenle CUTLASS, el yazısı PTX’in nasıl çalıştığını göstermek için harika bir örnektir.

GEMM, Top-K ve Softmax Birleşimi

Gösterdiğimiz belirli işlem, bir GEMM’in top_k ve softmax algoritmaları ile birleşimidir. Bu, bir uzman karışımı sinir ağı çalıştırırken sıkça kullanılan bir işlemdir. NVIDIA Hopper mimarisine odaklanacağız. Bu yaygın bir işlem olduğu için, CUTLASS’ın bu işlemi kapsayan özel bir çekirdeği mevcuttur ve bu, CUTLASS’ın yüksek performanslı GPU koduna el yazısı PTX’i nasıl entegre ettiğini göstermek açısından basit bir örnek olacaktır.

Bizim için şu anki koşullarda kullanılan:

• 3.9.2 numaralı CUTLASS sürümü
• NVIDIA GH200 GPU
• 570.140 sürüm numaralı sürücü
• 12.8 sürüm numaralı CUDA Toolkit

CUTLASS web sitesindeki derleme talimatlarını takip ederek, tam özellik kümesinin etkinleştirilmesi için -DCUTLASS_NVCC_ARCHS=90a derleme seçeneğini cmake ile kullanıyoruz. CUTLASS deposunda, en son mimariler üzerindeki çeşitli yetenekleri gösteren birçok örnek bulunmaktadır. CMake tamamlandıktan sonra, örnek kodun derlendiği geçen dizine (örneğin build/examples/61_hopper_gemm_with_topk_and_softmax) gidiyoruz. 

Yazılımı derlemek için make komutunu çalıştırdığımızda, kod derlenir ve çalıştırmaya hazır hale gelir. Uygulama, matris boyutları olan m, n, ve k ile hata toleransı epsilon ve benchmark sayılarının oluşturulması için çalıştırılacak yineleme sayısı gibi birkaç farklı girdi seçeneği alır.

Aşağıdaki çıktıyı, m=1024, n=8 (varsayılan), k=4096, iterations=1000000 ve epsilon=1e-4 değerlerini seçerek elde ederiz. Bu benchmarkta m token sayısını, n uzman sayısını, k uzmanların yerleştirme boyutunu gösterir ve top_k değeri test kodunda hard-coded olarak 2 olarak verilmiştir.

$ ./61_hopper_gemm_with_topk_and_softmax --m=1024 --k=4096  --iterations=1000000 --eps=1e-4
Sonuç: Geçti    Göreceli hata: 1.52478e-05
Problem Boyutu: 1024x8x4096x1
Ortalama çalışma süresi: 0.011765 ms
GFlop/s: 5704.11

Bu benchmark örneğinde, performans 5,704 GFlop/s olarak bulunmaktadır. Token sayısını (m parametresi) 16,384’e kadar çeşitlendirerek aşağıdaki performans tablosunu oluşturuyoruz.

Tablo 1. PTX’in kullanıldığı benchmark kodunun performansı, top_k ve softmax fonksiyonları

Inline PTX Kaldırma

Bu benchmark örneği, GEMM’i top_k ve softmax ile birleştirirken ve belirli koşullar altında inline PTX fonksiyonlarının kullanımıyla çalışıyor. Eğer k değeri 2 veya 4 ise, top_k fonksiyonu için inline PTX kullanıyor (not: Buradaki k matris boyutu k’sından farklıdır). Ayrıca, belirli koşullar altında softmax fonksiyonu için inline PTX kullanıyor. Her iki işlev için de belirli koşullar karşılanmadığında, CUDA C++ yazılmış yedekleme rutinleri devreye giriyor. Bu senaryoda PTX işlevlerinin kullanımıyla performans değişimini niceliklendirmek için , top_k ve softmax işlevlerindeki PTX çağrılarını yorum satırı haline getirmek yeterlidir.

Inline PTX’i bu örnekten çıkarmak için, cutlass/include/cutlass/epilogue/fusion/sm90_visitor_topk_softmax.hpp dosyasını açıp, inline PTX işlevlerinin kullanımını yorumlayarak kaldırıyoruz. Dosyanın en üst kısmında, top_2 ve top_4 ile başlayan inline PTX yazılmış bazı fonksiyonları göreceksiniz. İşte ilk PTX işlevinin örneği.

CUTLASS_DEVICE
Array<float, 2> top_2_reduce_scalar(Array<float, 2> a, float scalar) {
  Array<float, 2> out;
  asm volatile(
      "{n"
      "  .reg .f32 mx;n"
      "  .reg .pred p;n"
      "  max.f32 mx, %3, %4;n"
      "  setp.gtu.f32 p, %2, %4;n"
      "  selp.f32 %1, mx, %2, p;n"
      "  selp.f32 %0, %2, %4, p;n"
      "}n" : "=f"(out[0]), "=f"(out[1]) : "f"(a[0]), "f"(a[1]), "f"(scalar));
  return out;
}

Bunların tüm detaylarını anlamak zorunda değilsiniz. Burada asıl önemli olan, kısa bir inline PTX fonksiyonunun nasıl göründüğünü örneklemektir.

Aynı dosyada, bir softmax fonksiyonu da bulunmaktadır. Bunun gibi, çıkaracağımız PTX fonksiyonları işte bunlardır.

Ayrıca, o dosyada bu fonksiyonların if durumları içinde nasıl çağrıldığını görebilirsiniz. İnline PTX fonksiyonlarını çağıran if ifadelerini yorum satırı haline getirerek, yalnızca C++ implementasyonunun çalışma sürecine izin vermiş oluruz. 

Örneğin, add_element_to_desc_sorted_array adlı bir fonksiyon,  PTX fonksiyonları yerine, eğer k=2 veya k=4 ise, top_2 veya top_4 fonksiyonunu çağırabilir. İşlevin kodu:

void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& a, Element b) {
  if constexpr (N == 2 && is_same_v<Element, float>) {
    a = top_2_reduce_scalar(a, b);
  }
  else if constexpr (N == 4 && is_same_v<Element, float>) {
    a = top_4_reduce_scalar(a, b);
  }
  else {
    // daha yavaş genel yol, dallanma ile performans kaybı ve register sızıntısı oluşturabilir
    CUTLASS_PRAGMA_UNROLL
    for (int k = 0; k < N; ++k) {
      if (a[k] < b) {
        // Aşağı kaydır
        CUTLASS_PRAGMA_UNROLL
        for (int l = N - 1; l > k; --l) {
          a[l] = a[l-1];
        }
        a[k] = b;
        break;
      }
    }
  }
}

El yazısı PTX fonksiyonlarını çıkarmak için, top_2_reduce_scalar, top_4_reduce_scalar, top_2_reduce, top_4_reduce, ve fast_masked_softmax fonksiyonlarının çağrıldığı yerlerde aynı şekilde değişiklikler uygulanarak, PTX fonksiyonları yok edilir. İşte burada dikkat etmemiz gereken kod alanları:

void add_element_to_desc_sorted_array(cutlass::Array<Element, N>& a, Element b) {
/* BAŞLANGIÇ YORUMU
  if constexpr (N == 2 && is_same_v<Element, float>) {
    a = top_2_reduce_scalar(a, b);
  }
  else if constexpr (N == 4 && is_same_v<Element, float>) {
    a = top_4_reduce_scalar(a, b);
  }
  else {
 SON YORUM
    // daha yavaş genel yol, dallanma ile performans kaybı ve register sızıntısı oluşturabilir
    CUTLASS_PRAGMA_UNROLL
    for (int k = 0; k < N; ++k) {
      if (a[k] < b) {
        // Aşağı kaydır
        CUTLASS_PRAGMA_UNROLL
        for (int l = N - 1; l > k; --l) {
          a[l] = a[l-1];
        }
        a[k] = b;
        break;
      }
    }
//}  YORUMUN SONU
}

Tıpkı bu şekilde, merge_desc_sorted_arrays ve masked_softmax fonksiyonlarımızda da değişiklikler yaparak, el yazısı PTX işlevlerinin sıralı kullanımını ortadan kaldırıyoruz.

Aşağıdaki performans sonuçlarını elde ediyoruz.

Tablo 2. Yalnızca C++ kodu ile yazılan benchmark kodunun performansı

Bu sonuçları Tablo 1 ile karşılaştırdığımızda, el yazısı PTX kullanıldığında performansın %7 ila %14 arasında bir oranda arttığını görebiliyoruz. Buradaki önemli mesaj, bazı dikkatlice seçilmiş senaryolarda PTX yazarak performans elde etmenin mümkün olabileceğidir. El yazması PTX’in uygulamanıza dahil edilmesinin fizibilitesini belirlemek için performans ve taşınabilirlik takaslarının dikkatlice analiz edilmesi gerektiği sonucuna varıyoruz.

Bu, son derece optimize edilmiş bir örnek kodudur ve el yazısı PTX içeren bir senaryoyu göstermek için önemli bir örnek olarak seçilmiştir. 

Bu örnek, çoğu geliştiricinin PTX yazmayı bırakıp CUTLASS, CUBLAS ve diğer GPU kütüphanelerinin tasarımcılarına güvenmesi gerektiği yönündeki tavsiyeyi pekiştirmektedir.

Özet

Bu yazıda, CUTLASS’ın belirli bir birleşik GEMM işlemini geliştirmek için el yazısı PTX kullandığı bir örneği gösterdik. Her geliştiricinin PTX yazmasını gerektiği izlenimini vermek istemiyoruz. Çoğu geliştirici bu tür bir ihtiyaç hissetmeyecektir. El ile PTX yazımı, son çare olarak kullanılacak bir teknik olmalıdır.

Ancak, el yazısı PTX, tüm geliştiricilere açık bir tekniktir. Bu, ileri seviye ve özel bir tekniktir; uygun kullanıldığında, yüksek düzeyde GPU programcılarının araç kutusundaki bir başka araç olabilir.

CUDA platformunun büyük güçlerinden biri, geliştiricilerin her seviyede NVIDIA yığınından ihtiyaç duydukları şekilde faydalanabilmeleridir; uygulama seviyesinden başlayarak, montaj kodu yazmaya (PTX) kadar aşağıya kadar inebilirler.

Teşekkürler

Aşağıdaki NVIDIA katkıcısına teşekkür ederiz: Ali Hassani