“Warp 1.5.0’da Yer Tabanlı Programlama: Yeni Özellikler ve Kullanım Alanları”

Son çıkan Warp 1.5.0 ile geliştiricilerin Python’da yeni tile tabanlı programlama primitiflerine erişimi oldu. cuBLASDx ve cuFFTDx’ı kullanarak, bu yeni araçlar geliştiricilere etkin matris çarpımı ve Fourier dönüşümleri sunuyor. Bu blog yazısında, bu yeni özellikleri tanıtacak ve uygulamalarınızı optimize etmek için nasıl kullanılabileceklerini göstereceğiz. Warp 1.5.0’daki tile tabanlı programlama modeli şu anda ön izleme aşamasındadır; performans ve API’ler, gelecekteki sürümlerde değişiklik gösterebilir.

Giriş

Bugüne kadar, GPU donanımı, yalnızca SIMT (Single Instruction, Multiple Threads) yürütme modelinden, verimliliği artırmak için işbirlikçi işlemlere dayanan bir yapıya evrildi. Tensor Core matematik birimleri, genel GPU hesaplamalarının daha büyük bir parçası haline geldikçe, bu birimleri etkin bir şekilde ve üretken bir şekilde programlamak giderek daha önemli hale geliyor. BLAS gibi yüksek seviyeli API’ler, geniş bir yüksek performanslı düşük seviyeli talimat kümesine hedef verebilecek soyutlamalar sunmaktadır. Ancak, bu API’ler genellikle kullanıcı programlarıyla entegrasyonu zorlaştırmakta ve library çağrıları arasında sonuçların global bellekye geri alınmasını gerektirdiği için verimlilik kaybına neden olmaktadır. Öte yandan, Tensor Cores’u C++/CUDA seviyesinde doğrudan programlamak karmaşıktır ve veri akışının dikkatlice yönetimini gerektirir.

Bu sorunları ele almak amacıyla, OpenAI Triton ve C++ AMP gibi tile tabanlı programlama modelleri geliştirilmiştir. Saf SIMT modellerinin aksine, tile tabanlı programlama, geliştiricilerin birden fazla iş parçacığının işbirliği içinde gerçekleştirebileceği tile’lar üzerindeki işlemleri ifade etmelerine olanak tanır ve bu da verimliliği ve üretkenliği artırır.

Warp 1.5.0 sürümü ile Warp’ın çekirdek tabanlı programlama modelini tile tabanlı işlemlerle genişletiyoruz, bu da Warp geliştiricilerine modern GPU donanımının tüm gücüne erişim sağlamayı amaçlamaktadır. Bu uzantılar:

• SIMT’ten tile tabanlı yürütmeye sorunsuz geçiş sağlayan bir programlama modeli sunar.
• Elle indeksleme, paylaşılan bellek yönetimi ve gösterici aritmetiği gereksinimlerini azaltır.
• Geride yayılma ve eğitim için otomatik farklılaştırmayı destekler.

Ayrıca, Warp’ın cuBLASDx ve cuFFTDx’ı kullanarak, matris çarpımı ve Hızlı Fourier Dönüşümü (FFT) tile işlemlerini gerçekleştirmektedir. Warp’ın tile programlama modeli ile birleştirildiğinde, bu NVIDIA cihaz tarafı matematik kütüphaneleri, tek bir çekirdek içinde Tensor Core hızlandırmalı GEMM, FFT ve diğer tile işlemlerinin kesintisiz entegrasyonunu sağlar, böylece bellek I/O ve çekirdek başlatma overhead’ini azaltırken aritmetik yoğunluğu maksimize eder. Bu yaklaşım, yoğun lineer cebir gerektiren uygulamalar için geleneksel lineer cebir veya tensör çerçevelerine göre 4 kat daha iyi performans sağlamaktadır, örneğin robot ileri dinamiği uygulamalarında.

Warp Tile Primitifleri

Warp’taki yeni tile primitifleri, mevcut çekirdek tabanlı programlama modelini doğal bir şekilde genişletecek inşa etme, yükleme/saklama, lineer cebir ve harita/reduce işlemleri içerir.

İnşaat

Tile’lar, aşağıda gösterildiği gibi, Warp çekirdekleri içinde NumPy tarzı işlemler kullanılarak inşa edilebilir:

import warp as wp

@wp.kernel
def compute():
    # sıfırlı 32-bit float'lardan oluşan 16x16'lık bir tile oluştur
    a = wp.tile_zeros(m=16, n=16, dtype=wp.float32)
    # 1.0 ile başlatılmış 16-bit float'lardan oluşan 16x16'lık bir tile oluştur
    b = wp.tile_ones(m=16, n=16, dtype=wp.float16)

Warp’taki tile’lar, elemanları olarak skalar, vektör, matris veya yapılandırılmış veri tiplerini içeren iki boyutlu dizilerdir. Warp dizileri veya PyTorch tensörleri gibi, tile boyutları (örneğin yukarıdaki örnekteki 16×16 gibi), derleme zamanında bilinen sabitler olmalıdır. Ayrıca SIMT verilerinin bir iş parçacığına özgü olduğu gibi, tile verileri tüm CUDA bloğuna yayılır, ya register’larda ya da paylaşılan bellekte depolanır. Tile inşaat yöntemlerinin tam listesi burada GitHub’da bulunabilir.

Yükleme/Saklama

Warp, tile verilerini global bellekten yüklemek ve saklamak için açık yükleme/saklama işlemleri sağlar. Bu işlemler, tüm bloktaki tüm iş parçacıkları tarafından işbirlikçi şekilde gerçekleştirilir ve bu da global bellek ile paylaşılan veya register belleği arasında etkin veri aktarımını sağlar. Aşağıdaki örnekte, iki veri tile’ı global bellekten yüklenir, birlikte toplanır ve sonuç global belleğe geri saklanır. Kullanıcının paylaşılan bellek tahsisini veya saklamayı açıkça yönetmesi gerekmez:

import warp as wp

@wp.kernel
def compute(A: wp.array2d(dtype=float),
            B: wp.array2d(dtype=float),
            C: wp.array2d(dtype=float)):
    
    # girdi tile'larını işbirlikçi şekilde yükle
    a = wp.tile_load(A, i=0, j=0, m=16, n=16)
    b = wp.tile_load(B, i=0, j=0, m=16, n=16)

    # toplamı hesapla
    c = a + b

    # toplamı global belleğe işbirlikçi şekilde sakla
    wp.tile_store(C, i=0, j=0, t=c)

A = wp.ones((16,16), dtype=float)
B = wp.ones((16,16), dtype=float)
C = wp.empty((16,16), dtype=float)

wp.launch_tiled(compute, dim=1, inputs=[A, B, C], device="cuda:0", block_dim=64)

Yükleme/saklama işlemlerine ek olarak, Warp, wp.tile_atomic_add() gibi atomik işlemleri desteklemektedir. Bellek işlemlerinin tam listesi için lütfen belgelere göz atın.

Matris Çarpımı

Tile tabanlı programlamanın en önemli avantajlarından biri, işbirlikçi matris çarpımını gerçekleştirme yeteneğidir. Warp 1.5.0, geliştiricilerin işbirlikçi matris çarpımları gerçekleştirmelerini sağlayan wp.tile_matmul() adlı genel bir çarpma-toplama primitifini tanıtıyor. Bu primitif, eleman türlerine, matris boyutuna ve veri düzenine bağlı olarak en iyi performans için uygun Tensor Core MMA talimatlarını otomatik olarak yönlendiren cuBLASDx kullanılarak optimize edilmiştir.

Tile tabanlı programlama kullanarak Warp’ta matris çarpımı gerçekleştirmenin örneğini inceleyelim:

import warp as wp

TILE_M = wp.constant(32)
TILE_N = wp.constant(64)
TILE_K = wp.constant(64)

@wp.kernel
def gemm_tiled(A: wp.array2d(dtype=float), 
               B: wp.array2d(dtype=float), 
               C: wp.array2d(dtype=float)):

    i, j = wp.tid()

    # çıktı tile'ını ayır
    sum = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=float)
    count = int(K / TILE_K)

    # iç boyut üzerinde döngü
    for k in range(count):
        a = wp.tile_load(A, i, k, m=TILE_M, n=TILE_K)
        b = wp.tile_load(B, k, j, m=TILE_K, n=TILE_N)
 
        # matris çarpımı + biriktir
        wp.tile_matmul(a, b, sum)

    # sonucu sakla
    wp.tile_store(C, i, j, sum)

# 1024x1024 girdiler ile test et
M, N, K = 1024, 1024, 1024

A = wp.ones((M, K), dtype=float)
B = wp.ones((K, N), dtype=float)
C = wp.empty((M, N), dtype=float)

# 128 iş parçacığı ile çekirdeği başlat
wp.launch_tiled(gemm_tiled, 
                dim=(int(M//TILE_M), int(N//TILE_N)), 
                inputs=[A, B, C],
                block_dim=128)

Bu örnekte, gemm_tiled() adlı bir çekirdek tanımlıyoruz ve bir matris çarpımı işlemi gerçekleştiriyoruz. Çekirdek, global bellekten 2D veri dilimlerini yükler, bunları paylaşılan bellek tile’larına aktarır, wp.tile_matmul() kullanarak matris çarpımı yapar, sonucu paylaşılan bellekte biriktirir ve sonuçları global belleğe saklar.

Aşağıdaki şekilde, yukarıdaki GEMM çekirdeğinin performansını NVIDIA A100 80GB SXM’deki cuBLAS 12.4’e göre bir yüzde olarak görebilirsiniz (saat hızları maksimuma kilitlenmiştir) ve çeşitli FP32 matris boyutları için performans gözlemlenmektedir. Küçük problemler için, performansın cuBLAS ile rekabetçi olduğunu görmekteyiz; bu, optimal parametrelerin bu küçük boyut için bulunmuş olmasından ve başlatma overhead’inin maliyetin daha büyük bir kısmını oluşturmasından kaynaklanabilir. Ancak daha büyük problemler için performans daha düşük; çünkü şu anda tile sonuçlarının her zaman paylaşılan bellekte saklandığını görüyoruz. Yine basit bir örnek olmasına rağmen, büyük matrisler için cuBLAS performansının yaklaşık %70–80’i kadar bir performans sergilemektedir. Gelecek Warp sürümleri ve cuBLASDx, GEMM’lerin çıktısını register’larda tutarak performansı artırmayı sağlayacaktır.

Aşağıdaki şekilde, tek bir problem boyutunun genel performansı üzerinde tile boyutunun etkisini görebiliriz. Genel performans, problemin parçalanma biçimlerini belirleyen tile boyutlarına ve her alt probleme atanan iş parçacığı sayısını belirleyen blok boyutuna bağlıdır. Burada, M=N=K=1024 problemi için en iyi performansın TILE_M=32, TILE_N=64, TILE_K=64 ve 128 iş parçacığı kullanılmasıyla elde edildiğini görmekteyiz. Warp’ın dinamik programlama ve çalışma zamanı çekirdek oluşturma yetenekleri, kullanıcıların hyperparametre auto-tuning işlemlerini kolaylıkla gerçekleştirmelerine olanak tanımaktadır; bunu gösteren örnek için benchmark scriptine göz atın.

Lütfen referanse göz atarak tile lineer cebir primitiflerinin tam listesini bulabilirsiniz.

Map/Reduce İşlemleri

Warp 1.5.0 ayrıca, geliştiricilerin tile’lar üzerinde azaltmalar ve eleman bazlı işlemler gerçekleştirmelerini sağlayan map/reduce primitifleri içermektedir. Bu primitifler, LayerNorm ve SoftMax gibi işlemler için kritik öneme sahiptir.

Aşağıdaki örnek, bir dizilinin bir satırındaki tüm elemanların toplamını hesaplamak için her satır için bir CUDA bloğu kullanarak wp.tile_sum() ile işbirlikçi azaltma nasıl yapılacağını göstermektedir:

import warp as wp

@wp.kernel
def row_sum(input: wp.array2d(dtype=float),
            output: wp.array1d(dtype=float)):

    # blok indeksimizi al
    i = wp.tid()

    # global bellekten 256 elemanlık bir satırı yükle
    t = wp.tile_load(input[i], i=0, n=256)
    # elemanları işbirlikçi olarak topla
    s = wp.tile_sum(t)
    # toplamı çıktıya sakla
    wp.tile_store(output, i, s)

   
a = wp.ones((1024, 256), dtype=float)
b = wp.empty(1024, dtype=float)

wp.launch_tiled(row_sum, dim=[a.shape[0]], inputs=[a, b], block_dim=64)

Warp, özel azaltma operatörlerini desteklemektedir; bu örnekte wp.tile_reduce() ve wp.mul() kullanarak bir faktöriyel hesaplamaktayız, ancak kullanıcı tanımlı @wp.func azaltma operatörleri de kullanılabilir.

import warp as wp

@wp.kernel
def factorial():
    t = wp.tile_arange(1, 10, dtype=int)
    s = wp.tile_reduce(wp.mul, t)
    
    # "tile(m=1, n=1, storage=register) = [[362880]]" yazdırır
    print(s)

wp.launch(factorial, dim=[16], inputs=[], block_dim=16)

Map/reduce primitiflerinin tam listesi burada mevcuttur.

Vaka Çalışmaları

Sıkıştırılmış Sinir Ağları

Tile tabanlı programlama, sıkıştırılmış çok katmanlı algılayıcıların (MLP’ler) etkin bir şekilde uygulanmasını da sağlar. Aşağıda, Warp’taki tile tabanlı programlama kullanarak bir sıkıştırılmış MLP örneği verilmiştir:

import warp as wp

DIM_IN = wp.constant(4)
DIM_HID = wp.constant(32)
DIM_OUT = wp.constant(3)

@wp.kernel
def mlp_fused(weights_0: wp.array2d(dtype=wp.float16), 
              weights_1: wp.array2d(dtype=wp.float16), 
              loss: wp.array(dtype=float)):

    t = wp.tid()    

    # basit bir pozisyonel kodlama oluştur
    x = wp.vec4h(wp.sin(x), 
                 wp.cos(x),
                 wp.sin(x*2.0),
                 wp.cos(x*2.0))

    # özellik vektörleri oluşturmak için girişi tile'a yay
    f = wp.tile(x)
    
    # tamamen bağlı katman 0
    w0 = wp.tile_load(weights_0, 0, 0, m=DIM_HID, n=DIM_IN)
    z = wp.tile_map(relu, wp.tile_matmul(w0, f))

    # tamamen bağlı katman 1
    w1 = wp.tile_load(weights_1, 0, 0, m=DIM_OUT, n=DIM_HID)
    z = wp.tile_map(relu, wp.tile_matmul(w1, z))

    # kayıp fonksiyonu
    l = wp.tile_sum(z)

    wp.atomic_add(loss, 0, l)

wp.launch(mlp_fused, dim=(1,), inputs=[weights_0, weights_1, loss], block_dim=128)

Bu örnekte, mlp_fused() çekirdeği, ağırlıkları yükleyerek, matris çarpımları yaparak, wp.tile_map() kullanarak aktivasyon fonksiyonları uygulayarak ve kaybı hesaplayarak, tümünü tek bir çekirdek içinde gerçekleştirir. Aşağıdaki görüntü, bu yaklaşımı kullanarak bir görüntüyü kodlamak için bir örnek gösteriyor. Warp, otomatik farklılaştırmayı desteklediğinden, ağı ağırlıklarını değerlendirebilir ve görüntü koordinatlarından (x,y) piksel rengine (RGB) bir fonksiyon öğrenmek için eğitim alabilir. Tam örnek burada mevcuttur.

Signal İşleme

Warp tile işlemleri, cuFFTDx ile bütünleşerek, çekirdek içinde ileri ve ters FFT’ler sağlayarak, verilerin tile’ları üzerinde etkin Fourier dönüşümü işlemleri gerçekleştirir. İşte bir filtre kullanarak bir konvolüsyonu hesaplamak için Warp’ta tile tabanlı FFT kullanmanın bir örneği:

import warp as wp

@wp.kernel
def conv_tiled(x: wp.array2d(dtype=wp.vec2d), 
  y: wp.array2d(dtype=wp.vec2d), 
  z: wp.array2d(dtype=wp.vec2d)):
i, j = wp.tid()

# sinyali ve filtreyi yükle
a = wp.tile_load(x, i, j, m=TILE_M, n=TILE_N)
f = wp.tile_load(y, i, j, m=TILE_M, n=TILE_N)

# girdi sinyalinin Fourier dönüşümünü hesapla
wp.tile_fft(a)

# Frekans alanında filtreyi hesapla
c = wp.tile_map(cplx_prod, a, b)

# gerçek değerine geri dön
wp.tile_ifft(c)
wp.tile_store(z, i, j, c)

Bu örnekte, conv_tiled() çekirdeği, bir tile verisinin ileri FFT’sini (son boyutta) gerçekleştirir, bir filtre uygular ve ardından ters FFT hesaplar. Uygulama altında, cuFFTDx kullanılmaktadır. Tam örnek için burada ulaşabilirsiniz. Aşağıdaki görüntü, gürültülü bir girdi sinyaline uygulanan filtrenin çıktısını göstermektedir.

Robot İleri Dinamikleri

Tile tabanlı programlama, yoğun lineer cebir gerektiren simülasyon uygulamaları için de oldukça faydalıdır. Robot simülasyonunda, Composite Rigid Body Algorithm (CRBA) yöntemi, eklem mekanizmaları için ileri dinamiklerini hesaplamak için kullanılır. CRBA yönteminde, aşağıdaki üçlü-matris çarpımı gereklidir, burada içteki M bir blok seyrek diyagonal kütle matrisidir:

Oluşturulduktan sonra, sistem matrisi Cholesky ayrıştırması kullanılarak çözülür ve ileri ve geri yerine koyma ile çözülür. Bu problemi, Warp’ın tile primitiflerini kullanarak seyrekliğin avantajını alacak şekilde bir batch sürümünde ifade edebiliriz:

import warp as wp

@wp.kernel
def forward_dynamics(
    J_arr: wp.array3d(dtype=float),
    M_arr: wp.array3d(dtype=float),
    R_arr: wp.array3d(dtype=float),
    H_arr: wp.array3d(dtype=float),
    L_arr: wp.array3d(dtype=float),
):
    batch = wp.tid()

    J = wp.tile_load(J_arr[batch], 0, 0, 
  m=wp.static(6 * num_joints), n=num_dofs)
    P = wp.tile_zeros(m=wp.static(6 * num_joints), n=num_dofs, dtype=float)

    # P = M*J hesabı yapın, burada M 6x6 blok diyagonel kütle matrisidir
    for i in range(int(num_joints)):

        # 6x6 blok matrisler diyagonal üzerinde
        M_body = wp.tile_load(M_arr[batch], i, i, m=6, n=6)

        # Jacobian'dan bir 6xN satırı yükle
        J_body = wp.tile_view(J, i * 6, 0, m=6, n=num_dofs)

        # ağırlıklı satırı hesapla
        P_body = wp.tile_matmul(M_body, J_body)

        # P dilimine atama yap
        wp.tile_assign(P, i * 6, 0, P_body)

    # H = J^T*P hesapla
    H = wp.tile_matmul(wp.tile_transpose(J), P)

    # cholesky L L^T = (H + diag(R))
    R = wp.tile_load(R_arr[batch], 0, 0, m=num_dofs, n=1, storage="shared")

    H += wp.tile_diag(R)
    L = wp.tile_cholesky(H)

    wp.tile_store(L_arr[batch], 0, 0, L)

# 64 iş parçacığı ile robot başına çekirdeği başlat
wp.launch_tiled(forward_dynamics, 
                dim=(num_robots,), 
                inputs=[J_arr, M_arr, R_arr, H_arr, L_arr], 
                block_dim=64)

Bu örnekte, forward_dynamics() çekirdeği, CRBA yöntemini, Jacobian ve kütle matrislerinin tile’larını yükleyerek gerçekleştirir ve ürünlerini H sistem matrisini ve Cholesky faktörizasyonunu oluşturmak için hesaplar. Torch, bu özel kullanım durumunda bir düzine çekirdek başlatmak gerektirirken, Warp uygulaması tek bir tam birleşik çekirdek gerektirir. Bu, global bellek dönüşlerinin ve başlatma overhead’inin miktarını azaltarak önemli ölçüde daha iyi performans sağlamaktadır.

1,024 dört ayaklı robot için performans aşağıdaki gibidir; ileri dinamikler çekirdeği bir NVIDIA A100 80GB GPU’da çalışırken (tüm zamanlamalar milisaniye cinsindendir); daha düşük değer daha iyidir:

Tam örnek burada mevcut olup, Cholesky ayrıştırmalarını ve geri yerine koymayı da içeren uzantıların yakında geleceğini belirtmeliyim.

Gelecek Gelişmeler

Gelecek Warp ve MathDx sürümleri, aşağıdaki özellikleri içerecektir:

• Satır bazlı azaltma operatörleri için ek destek
• Lambda fonksiyonlarından tile oluşturma
• Veri tipi ve düzen dönüşümleri
• GEMM operasyonları için geliştirilmiş performans
• Çeşitli matris ayrıştırma algoritmaları dahil olmak üzere ek lineer cebir primitifleri.

Daha Fazla Bilgi Edinin

Warp 1.5.0’daki tile tabanlı programlama, geliştiricilerin uygulamalarında önemli performans iyileştirmeleri sağlamak için güçlü ve esnek bir yaklaşım sunmaktadır. cuBLASDx ve cuFFTDx’ı kullanarak, Warp 1.5.0, GEMM ve FFT işlemlerinin kesintisiz birleşimini sağlayarak bellek I/O ve çekirdek başlatma overhead’ini azaltmaktadır.

Warp’ın Tile işlemlerini kullanmaya başlamak için şu komut ile Python ortamınıza Warp’ı yükleyin:

pip install warp-lang

Sıkıştırılmış MLP örneğini çalıştırmak için aşağıdaki komutu kullanın:

python -m warp.examples.tile.example_mlp.py

İlgili Kaynaklar

Warp 1.5.0 ve NVIDIA Matematik cihaz hızlandırma (Dx) kütüphaneleri hakkında daha fazla bilgi edinmek için aşağıdaki bağlantılara göz atabilirsiniz:

• Warp GTC Sunumu: https://www.nvidia.com/en-us/on-demand/session/gtc24-s63345/
• Warp Kaynağı: https://github.com/nvidia/warp
• Warp Tile Belgeleri: https://nvidia.github.io/warp/modules/tiles.html
• cuBLASDx: https://developer.nvidia.com/cublasdx-downloads
• cuFFTDx: https://developer.nvidia.com/cufftdx-downloads

Teşekkürler

Bu yazıya ve projeye katkıda bulunan Paweł Grabowski, Doris Pan, Neil Lindquist, Jakub Szuppe, Łukasz Ligowski, Sergey Maydanov ve Łukasz Wawrzyniak’a teşekkür ederiz.