Bentuk dan tata letak

Struktur Operasi XLA

Pertimbangkan contoh HLO:

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

Hal ini terdiri dari komponen berikut:

  • Nama Op: add.936
    • Ini adalah nama unik untuk operasi.
  • Bentuk: bf16[8,1,1280,16384]
    • Ini adalah bentuk output Op. Di sini, dtype-nya adalah bf16 dan bentuknya adalah [8,1,1280,16384].
  • Tata Letak (dengan Pengubinan): 3,2,0,1:T(8,128)(2,1)
    • Hal ini menjelaskan cara array disimpan dalam memori. 3,2,0,1 menunjukkan urutan sumbu dalam memori (yaitu, kolom utama, baris utama, dll.) dan T(8,128)(2,1) menunjukkan pengelompokan & padding yang digunakan.
    • Tata letak bersifat opsional. Jika tidak ditentukan, tidak ada pengubinan dan dimensi diasumsikan diurutkan dari yang paling utama hingga yang paling kecil.
  • Operasi: add
    • Operasi yang sedang dilakukan. Di sini, nilainya adalah Add, yang juga disebutkan dalam nama Op.
  • Argumen: exponential.183, broadcast.3115
    • Operasi ini menggunakan dua argumen, yang ditentukan dengan nama uniknya.

Mari kita lihat contoh lain, yaitu operasi gabungan:

%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

Selain komponen yang dijelaskan sebelumnya, komponen ini terdiri dari:

  • Atribut: kind dan calls
    • Objek ini memberikan informasi selengkapnya tentang operasi yang sedang dilakukan, dalam hal ini: fusi.
  • Lokasi memori (ID ruang memori): S(1)
    • Ini menunjukkan ruang/lokasi memori tempat array disimpan. S(1) di sini menunjukkan bahwa array ini berada di VMEM (di TPU).
  • Detail bentuk dan tata letak untuk argumen input %fusion.32

Bagian berikut menjelaskan Bentuk, Tata Letak, dan ID Ruang Memori. Anda dapat mempelajari lebih lanjut Penataan dalam Tata Letak Berpetak.

Bentuk

Proto XLA ShapeProto(xla_data.proto) mendeskripsikan jumlah dimensi, ukuran, dan jenis data array N-dimensi (singkatnya array).

Terminologi, notasi, dan konvensi

  • Jumlah dimensi sebenarnya array adalah jumlah dimensi yang memiliki ukuran lebih besar dari 1.

  • Dimensi diberi nomor dari 0 hingga N-1 untuk array berdimensi N. Ukuran dimensi adalah bilangan bulat non-negatif. Khususnya, ukuran 0 valid. Nomor dimensi adalah label arbitrer untuk memudahkan. Urutan angka dimensi ini tidak menyiratkan urutan minor/mayor tertentu dalam tata letak bentuk. Tata letak ditentukan oleh proto LayoutProto.

  • Menurut konvensi, dimensi dicantumkan dalam urutan menaik dari nomor dimensi. Misalnya, untuk array 3 dimensi berukuran [A x B x C], dimensi 0 berukuran A, dimensi 1 berukuran B, dan dimensi 2 berukuran C.

    Beberapa utilitas di XLA juga mendukung pengindeksan negatif seperti Python: Dimensi -1 adalah dimensi terakhir (setara dengan N-1 untuk array berdimensi N). Misalnya, untuk array 3 dimensi yang dijelaskan di atas, dimensi -1 memiliki ukuran C, dimensi -2 memiliki ukuran B, dan seterusnya.

  • Array dua, tiga, dan empat dimensi sering kali memiliki huruf tertentu yang terkait dengan dimensi. Misalnya, untuk array 2D:

    • dimensi 0: y
    • dimensi 1: x

    Untuk array 3D:

    • dimensi 0: z
    • dimensi 1: y
    • dimensi 2: x

    Untuk array 4D:

    • dimensi 0: p
    • dimensi 1: z
    • dimensi 2: y
    • dimensi 3: x

  • Fungsi di XLA API yang menggunakan dimensi melakukannya dalam urutan peningkatan jumlah dimensi. Ini cocok dengan pengurutan yang digunakan saat meneruskan dimensi sebagai initializer_list; misalnya

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

    akan membuat bentuk yang array ukuran dimensinya terdiri dari urutan [A, B, C, D].

Tata Letak

Proto LayoutProto menjelaskan cara array direpresentasikan dalam memori. Kolom ini mencakup kolom berikut:

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

Pengurutan dimensi dari kecil ke besar

Satu-satunya kolom yang wajib diisi adalah minor_to_major. Kolom ini menjelaskan pengurutan dimensi dari kecil ke besar dalam bentuk. Nilai dalam minor_to_major adalah pengurutan dimensi array (0 hingga N-1 untuk array berdimensi N) dengan nilai pertama menjadi dimensi paling kecil hingga nilai terakhir yang merupakan dimensi paling besar. Dimensi yang paling kecil adalah dimensi yang berubah paling cepat saat menelusuri elemen array yang disusun dalam memori linear.

Misalnya, pertimbangkan array 2D berukuran [2 x 3] berikut:

a b c
d e f

Di sini, dimensi 0 adalah ukuran 2, dan dimensi 1 adalah ukuran 3. Jika kolom minor_to_major dalam tata letak adalah [0, 1], maka dimensi 0 adalah dimensi paling kecil dan dimensi 1 adalah dimensi paling besar. Hal ini sesuai dengan tata letak berikut dalam memori linear:

a d b e c f

Urutan dimensi kecil ke besar dari 0 hingga N-1 mirip dengan column-major (untuk 2 dimensi). Dengan asumsi pengurutan dimensi yang monoton, cara lain kita dapat merujuk pada tata letak ini dalam kode hanyalah "dim 0 adalah minor".

Di sisi lain, jika kolom minor_to_major dalam tata letak adalah [1, 0], maka tata letak dalam memori linear adalah:

a b c d e f

Urutan dimensi minor ke mayor dari N-1 hingga 0 untuk array berdimensi N mirip dengan row-major (untuk 2 dimensi). Dengan asumsi pengurutan dimensi yang monoton, cara lain untuk merujuk tata letak ini dalam kode adalah dengan "dimensi 0 adalah dimensi utama".

Urutan default dari versi minor ke versi mayor

Tata letak default untuk Bentuk yang baru dibuat adalah "dimension order is major-to-minor" (yaitu [N-1, ..., 0]).

Padding

Kolom tail_padding_alignment_in_elements menentukan perataan array berpetak dalam hal jumlah elemen. Setelah menerapkan pengelompokan, elemen yang diberi padding akan ditambahkan di akhir tata letak hingga jumlah total elemen merupakan kelipatan dari nilai ini.

Mengindeks ke dalam array

Class IndexUtil di index_util.h menyediakan utilitas untuk mengonversi antara indeks multidimensi dan indeks linier berdasarkan bentuk dan tata letak. Indeks multidimensi mencakup indeks int64 untuk setiap dimensi. Indeks linear adalah nilai int64 tunggal yang mengindeks ke dalam buffer yang menyimpan array. Lihat shape_util.h dan layout_util.h di direktori yang sama untuk utilitas yang menyederhanakan pembuatan dan manipulasi bentuk dan tata letak.

ID Ruang Memori

Di HLO, setiap array dapat diberi anotasi dengan ID ruang memori, yang ditulis sebagai S(n).

  • S(0) (sering dihilangkan) menunjukkan memori bandwidth tinggi (HBM) perangkat.
  • S(1) mewakili memori virtual (VMEM) di perangkat.
  • S(2), S(3), dll., sesuai dengan ruang memori tambahan khusus perangkat.
  • S(5) menunjukkan memori host.