Şekiller ve düzen

XLA İşleminin Yapısı

Örnek bir HLO'yu inceleyelim:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

Bu, aşağıdaki bileşenlerden oluşur:

  • İşlem Adı: add.936
    • Bu, işlemin benzersiz adıdır.
  • Şekil: bf16[8,1,1280,16384]
    • Bu, Op'nin çıkış şeklidir. Burada dtype, bf16 ve şekil [8,1,1280,16384]'dir.
  • Düzen (döşeme ile): 3,2,0,1:T(8,128)(2,1)
    • Bu, dizinin bellekte nasıl depolandığını açıklar. 3,2,0,1, bellekteki eksenlerin sırasını (ör.sütun öncelikli, satır öncelikli vb.) belirtir. T(8,128)(2,1) ise kullanılan döşeme ve dolguyu belirtir.
    • Düzen isteğe bağlıdır. Belirtilmezse döşeme yapılmaz ve boyutların en büyükten en küçüğe doğru sıralandığı varsayılır.
  • İşlem: add
    • Gerçekleştirilen işlem. Burada, Op adında da belirtilen Add (Ekle) değeri kullanılır.
  • Arguments: exponential.183, broadcast.3115
    • Bu işlem, benzersiz adlarıyla belirtilen iki bağımsız değişken alır.

Şimdi de bir füzyon işlemi örneğini inceleyelim:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

Daha önce açıklanan bileşenlere ek olarak şunlardan oluşur:

  • Özellikler: kind ve calls
    • Bu parametreler, gerçekleştirilen işlem (bu örnekte birleştirme) hakkında daha fazla bilgi sağlar.
  • Bellek konumu (bellek alanı tanımlayıcısı): S(1)
    • Bu, dizinin depolandığı bellek alanını/konumunu gösterir. S(1) burada, bu dizinin VMEM'de (TPU'da) bulunduğunu gösterir.
  • %fusion.32 giriş bağımsız değişkeninin şekil ve düzen ayrıntıları

Aşağıdaki bölümlerde Şekiller, Düzen ve Bellek Alanı Tanımlayıcıları açıklanmaktadır. Döşeme hakkında daha fazla bilgiyi Döşeme Düzeni başlıklı makalede bulabilirsiniz.

Şekiller

XLA ShapeProto proto (xla_data.proto) N boyutlu bir dizinin (kısaca dizi) boyut sayısını, boyutunu ve veri türünü açıklar.

Terminoloji, notasyon ve kurallar

  • Bir dizinin gerçek boyut sayısı, boyutu 1'den büyük olan boyutların sayısıdır.

  • Boyutlar, N boyutlu bir dizi için 0 ile N-1 arasında numaralandırılır. Bir boyutun boyutu, negatif olmayan bir tam sayıdır. Özellikle 0 boyutu geçerlidir. Boyut numaraları, kolaylık sağlamak için rastgele etiketlerdir. Bu boyut numaralarının sırası, şeklin düzeninde belirli bir küçük/büyük sıralama anlamına gelmez. Düzen, LayoutProto proto'su tarafından belirlenir.

  • Boyutlar, boyut numarasına göre artan düzende listelenir. Örneğin, [A x B x C] boyutunda 3 boyutlu bir dizi için boyut 0'ın boyutu A, boyut 1'in boyutu B ve boyut 2'nin boyutu C'tür.

    XLA'daki bazı yardımcı programlar, Python benzeri negatif indekslemeyi de destekler: Boyut -1, son boyuttur (N boyutlu bir dizi için N-1'ya eşdeğerdir). Örneğin, yukarıda açıklanan 3 boyutlu dizi için -1 boyutu C boyutuna, -2 boyutu B boyutuna sahiptir.

  • İki, üç ve dört boyutlu dizilerde genellikle boyutlarla ilişkili belirli harfler bulunur. Örneğin, 2 boyutlu bir dizi için:

    • 0. boyut: y
    • 1. boyut: x

    3D dizi için:

    • 0. boyut: z
    • 1. boyut: y
    • 2. boyut: x

    4 boyutlu bir dizi için:

    • 0. boyut: p
    • 1. boyut: z
    • 2. boyut: y
    • 3. boyut: x

  • XLA API'deki boyut alan işlevler, boyut numarası artan sırada çalışır. Bu, boyutlar initializer_list olarak iletilirken kullanılan sıralamayla eşleşir. Örneğin:

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    boyut dizisi [A, B, C, D] olan bir şekil oluşturur.

Düzen

LayoutProto proto'su, bir dizinin bellekte nasıl temsil edildiğini açıklar. Aşağıdaki alanları içerir:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

Küçükten büyüğe boyut sıralaması

Tek zorunlu alan minor_to_major'dır. Bu alan, bir şekil içindeki boyutların küçükten büyüğe sıralanmasını açıklar. minor_to_major içindeki değerler, dizinin boyutlarının sıralamasıdır (N boyutlu bir dizi için 0 ile N-1 arasında). İlk değer en küçük boyuttan başlayıp son değer olan en büyük boyuta kadar devam eder. En küçük boyut, doğrusal bellekte düzenlenmiş dizinin öğeleri arasında ilerlerken en hızlı değişen boyuttur.

Örneğin, [2 x 3] boyutundaki aşağıdaki 2 boyutlu diziyi ele alalım:

a b c
d e f

Burada 0 boyutu 2, 1 boyutu ise 3'tür. Düzendeki minor_to_major alanı [0, 1] ise 0 boyutu en küçük boyut, 1 boyutu ise en büyük boyuttur. Bu, doğrusal bellekteki aşağıdaki düzene karşılık gelir:

a d b e c f

0 ile N-1 arasındaki bu küçükten büyüğe boyut sırası, sütun öncelikli (2 boyutlu diziler için) sıraya benzer. Boyutların monoton bir şekilde sıralandığını varsayarsak bu düzenden kodda "dim 0 is minor" (0. boyut küçük) şeklinde de bahsedebiliriz.

Öte yandan, düzendeki minor_to_major alanı [1, 0] ise doğrusal bellekteki düzen şu şekildedir:

a b c d e f

N boyutlu bir dizide N-1 ile 0 arasında değişen küçükten büyüğe boyut sırası, satır öncelikli (2 boyutlu diziler için) sıraya benzer. Boyutların monoton bir şekilde sıralandığını varsayarsak bu düzene kodda "0. boyut ana boyuttur" şeklinde de atıfta bulunabiliriz.

Varsayılan olarak küçükten büyüğe sıralama

Yeni oluşturulan şekillerin varsayılan düzeni "boyut sırası büyükten küçüğe" şeklindedir (ör. [N-1, ..., 0]).

Dolgu

tail_padding_alignment_in_elements alanı, döşeli dizinin öğe sayısı açısından hizalamasını tanımlar. Döşeme uygulandıktan sonra, toplam öğe sayısı bu değerin katı olana kadar düzenin sonuna doldurulmuş öğeler eklenir.

Dizilere dizin oluşturma

index_util.h dosyasındaki IndexUtil sınıfı, şekil ve düzen verildiğinde çok boyutlu dizinler ile doğrusal dizinler arasında dönüştürme için yardımcı programlar sağlar. Çok boyutlu dizinler, her boyut için bir int64 dizini içerir. Doğrusal dizinler, diziyi tutan arabellekte dizin oluşturmak için kullanılan tek bir int64 değeridir. Şekillerin ve düzenlerin oluşturulmasını ve değiştirilmesini kolaylaştıran yardımcı programlar için aynı dizindeki shape_util.h ve layout_util.h bölümlerine bakın.

Bellek alanı tanımlayıcıları

HLO'da her dizi, S(n) olarak yazılan bir bellek alanı tanımlayıcısıyla açıklama eklenmiş olabilir.

  • S(0) (genellikle atlanır), cihazın yüksek bant genişlikli belleğini (HBM) ifade eder.
  • S(1), cihazdaki sanal belleği (VMEM) temsil eder.
  • S(2), S(3) vb., cihaza özel ek bellek alanlarına karşılık gelir.
  • S(5), ana makine belleğini gösterir.