Jenis Tensor Sparse di MLIR

Encoding tensor renggang adalah atribut untuk mengenkode informasi pada properti ketersebaran tensor, yang terinspirasi oleh formalisasi TACO dari tensor sparse. Pengkodean ini akhirnya digunakan dengan penerusan sparsifier untuk menghasilkan kode sparse secara otomatis dari representasi agnostik ketersebaran dari komputasi, yaitu, representasi sparse implisit dikonversi menjadi representasi renggang eksplisit di mana loop iterasi bersama beroperasi pada penyimpanan sparse daripada tensor dengan ketersebaran encoding. Penerusan compiler yang berjalan sebelum penerusan sparsifier ini perlu diketahui semantik tipe tensor dengan encoding ketersebaran.

Dalam encoding ini, kami menggunakan dimensi untuk lihat sumbu tensor semantik, dan level untuk merujuk ke sumbu format penyimpanan yang sebenarnya, yaitu, representasi operasional dari tensor sparse di memori. Jumlah dimensi data biasanya sama dengan jumlah tingkat (misalnya format penyimpanan CSR). Namun, pengkodean juga dapat memetakan dimensi ke tingkat yang lebih tinggi (misalnya, untuk mengenkode format penyimpanan BSR blok-sparse) atau ke tingkat urutan yang lebih rendah (misalnya, untuk linearisasi dimensi sebagai satu tingkat dalam penyimpanan).

Encoding berisi peta yang memberikan hal berikut:

  • Urutan spesifikasi dimensi yang diurutkan, yang masing-masing menentukan:
    • dimensi-size (implisit dari bentuk dimensi tensor)
    • ekspresi dimensi
  • Urutan spesifikasi level yang berurutan, yang masing-masing mencakup level-type, yang menentukan cara level disimpan. Setiap jenis level terdiri dari:
    • level-expression, yang menentukan apa yang disimpan
    • format-level
    • kumpulan properti tingkat yang berlaku untuk format tingkat

Setiap ekspresi tingkat adalah ekspresi affine terhadap variabel dimensi. Dengan demikian, ekspresi-tingkat secara kolektif mendefinisikan peta affine dari koordinat dimensi ke koordinat level. Ekspresi dimensi secara kolektif menentukan peta terbalik, yang hanya perlu diberikan untuk kasus yang rumit yang tidak dapat disimpulkan secara otomatis.

Setiap dimensi juga dapat memiliki SparseTensorDimSliceAttr opsional. Dalam format penyimpanan sparse, kita mengacu pada indeks yang disimpan secara eksplisit sebagai koordinat dan offset ke dalam format penyimpanan sebagai posisi.

Format tingkat yang didukung adalah sebagai berikut:

  • dense : semua entri di sepanjang level ini disimpan
  • dikompresi : hanya angka bukan nol di sepanjang level ini yang disimpan
  • loose_compressed : saat dikompresi, tetapi memungkinkan adanya ruang kosong antar-region
  • singleton : varian dari format terkompresi, yang koordinatnya tidak memiliki saudara
  • block2_4 : kompresi menggunakan encoding 2:4 per blok 1x4

Untuk level terkompresi, setiap interval posisi dinyatakan dalam diagram arah dengan batas bawah pos(i) dan batas atas pos(i+1) - 1, yang menyiratkan bahwa interval yang berurutan harus muncul secara berurutan tanpa ada "lubang" di antara mereka. Format kompresi longgar melonggarkan batasan ini dengan merepresentasikan interval posisi dengan batas bawah lo(i) dan batas atas hi(i), yang memungkinkan interval untuk muncul dalam urutan sewenang-wenang dan dengan ruang siku di antara mereka.

Secara default, setiap jenis tingkat memiliki properti unik (tidak ada duplikat koordinat di tingkat itu) dan diurutkan (koordinat muncul diurutkan pada tingkat itu tertentu). Properti berikut dapat ditambahkan ke format tingkat untuk mengubah perilaku default ini:

  • nonunik : koordinat duplikat mungkin muncul pada tingkat
  • tidak berurutan : koordinat dapat ditampilkan dalam urutan arbribratri

Selain peta, dua kolom berikut bersifat opsional:

  • Bitwidth yang diperlukan untuk penyimpanan posisi (offset integral ke dalam skema penyimpanan sparse). Lebar yang sempit mengurangi memori jejak penyimpanan overhead, selama lebarnya sudah cukup untuk menentukan total rentang yang diperlukan (yaitu jumlah maksimum entri di semua tingkat pengalihan tidak langsung). Pilihannya adalah 8, 16, 32, 64, atau, default, 0 untuk menunjukkan bitwidth native.

  • Bitwidth yang diperlukan untuk penyimpanan koordinat (koordinat entri yang disimpan). Lebar yang sempit akan mengurangi jejak memori penyimpanan {i>overhead<i}, selama lebarnya sudah memadai untuk total rentang yang diperlukan (yaitu nilai maksimum setiap tensor berkoordinasi di semua tingkatan). Pilihannya adalah 8, 16, 32, 64, atau, default, 0 untuk menunjukkan bitwidth native.

Contoh

Dalam format CSR(Compressed Sparse Row) seperti ditunjukkan di bawah ini, adalah:

#CSR = #sparse_tensor.encoding<{
  map = (i, j) -> (i : dense, j : compressed)
}>

Hal ini menunjukkan bahwa first dimension (baris) dipetakan ke first level, yang merupakan level dense, ditunjukkan dengan ukuran 4. Dan second dimension (kolom) dipetakan ke second level, yang ditunjukkan oleh array posisi dan koordinat objek. Nilai 3 ([1, 1] dalam matriks asli) adalah direpresentasikan oleh offset dari array posisi (nilai nomor baris 3 adalah 1 dalam matriks asli karena merupakan pasangan offset kedua dan nomor kolomnya berada di indeks [2 : 4) pada array koordinat). Pada bagian , kita dapat melihat nilai itu untuk nomor kolom 3 adalah 1 di matriks asli.

Untuk format BSR(Block Sparse Row), jenis tensor sparse adalah:

#BSR = #sparse_tensor.encoding<{
  map = (i, j) ->
    ( i floordiv 2 : dense
    , j floordiv 2 : compressed
    , i mod 2 : dense
    , j mod 2 : dense
    )

Pertimbangkan matriks sparse berikut, dengan blok 2x2:

Example 2x2 block storage:
 +-----+-----+-----+    +-----+-----+-----+
 | 1 2 | . . | 4 . |    | 1 2 |     | 4 0 |
 | . 3 | . . | . 5 |    | 0 3 |     | 0 5 |
 +-----+-----+-----+ => +-----+-----+-----+
 | . . | 6 7 | . . |    |     | 6 7 |     |
 | . . | 8 . | . . |    |     | 8 0 |     |
 +-----+-----+-----+    +-----+-----+-----+

yang pada akhirnya disimpan dalam format sesuai TACO sebagai

Stored as:
   positions[1]   : 0 2 3
   coordinates[1] : 0 2 1
   values         : 1.000000 2.000000 0.000000 3.000000
                    4.000000 0.000000 0.000000 5.000000
                    6.000000 7.000000 8.000000 0.000000

Yang, kebetulan, merupakan Format blok NVidia yang disajikan dalam dokumentasi cuSparse.

Blokir baris sparse. (n.d.-b). Nvidia. https://docs.nvidia.com/cuda/cusparse/_images/bsr.png

Kami juga mendukung Nvidia's 2:4 structured sparsity.

Jenis tensor sparse untuk ini adalah sebagai berikut:

#NV_24 = #sparse_tensor.encoding<{
  map = ( i, j ) -> ( i            : dense,
                      j floordiv 4 : dense,
                      j mod 4      : block2_4),
  crdWidth = 2  // 2-bits for each coordinate
}>

Berdasarkan matriks sampel yang diberikan dalam dokumentasi NVidia:

Contoh penyimpanan MMA sparse. (tanpa tahun).: Nvidia. https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-sparse-mma

MLIR sekarang memetakan matriks ini ke tata letak yang identik:

coordinates[2]  :
   0 2 0 2 0 2 0 2
   1 3 1 3 1 3 1 3
   0 1 2 3 0 1 2 3
   2 3 0 1 2 3 0 1
   0 1 0 1 0 1 0 1
   0 1 0 1 0 1 0 1
   2 3 2 3 2 3 2 3
   2 3 2 3 2 3 2 3
   0 2 0 2 0 2 0 2
   1 3 1 3 1 3 1 3
   0 1 2 3 0 1 2 3
   2 3 0 1 2 3 0 1
   0 1 0 1 0 1 0 1
   0 1 0 1 0 1 0 1
   2 3 2 3 2 3 2 3
   2 3 2 3 2 3 2 3
values :
  1.000000 2.000000 3.000000 4.000000 1.000000 2.000000 3.000000 4.000000
  5.000000 6.000000 7.000000 8.000000 5.000000 6.000000 7.000000 8.000000
  9.000000 10.000000 11.000000 12.000000 9.000000 10.000000 11.000000 12.000000
  13.000000 14.000000 15.000000 16.000000 13.000000 14.000000 15.000000 16.000000
  17.000000 18.000000 19.000000 20.000000 17.000000 18.000000 19.000000 20.000000
  21.000000 22.000000 23.000000 24.000000 21.000000 22.000000 23.000000 24.000000
  25.000000 26.000000 27.000000 28.000000 25.000000 26.000000 27.000000 28.000000
  29.000000 30.000000 31.000000 32.000000 29.000000 30.000000 31.000000 32.000000
  1.000000 2.000000 3.000000 4.000000 1.000000 2.000000 3.000000 4.000000
  5.000000 6.000000 7.000000 8.000000 5.000000 6.000000 7.000000 8.000000
  9.000000 10.000000 11.000000 12.000000 9.000000 10.000000 11.000000 12.000000
  13.000000 14.000000 15.000000 16.000000 13.000000 14.000000 15.000000 16.000000
  17.000000 18.000000 19.000000 20.000000 17.000000 18.000000 19.000000 20.000000
  21.000000 22.000000 23.000000 24.000000 21.000000 22.000000 23.000000 24.000000
  25.000000 26.000000 27.000000 28.000000 25.000000 26.000000 27.000000 28.000000
  29.000000 30.000000 31.000000 32.000000 29.000000 30.000000 31.000000 32.000000

Contoh lainnya:

// Sparse vector.
#SparseVector = #sparse_tensor.encoding<{
  map = (i) -> (i : compressed)
}>
... tensor<?xf32, #SparseVector> ...

// Sorted coordinate scheme.
#SortedCOO = #sparse_tensor.encoding<{
  map = (i, j) -> (i : compressed(nonunique), j : singleton)
}>
... tensor<?x?xf64, #SortedCOO> ...

// Batched sorted coordinate scheme, with high encoding.
#BCOO = #sparse_tensor.encoding<{
  map = (i, j, k) -> (i : dense, j : compressed(nonunique, high), k : singleton)
}>
... tensor<10x10xf32, #BCOO> ...

// Compressed sparse row.
#CSR = #sparse_tensor.encoding<{
  map = (i, j) -> (i : dense, j : compressed)
}>
... tensor<100x100xbf16, #CSR> ...

// Doubly compressed sparse column storage with specific bitwidths.
#DCSC = #sparse_tensor.encoding<{
  map = (i, j) -> (j : compressed, i : compressed),
  posWidth = 32,
  crdWidth = 8
}>
... tensor<8x8xf64, #DCSC> ...

// Block sparse row storage (2x3 blocks).
#BSR = #sparse_tensor.encoding<{
  map = ( i, j ) ->
  ( i floordiv 2 : dense,
    j floordiv 3 : compressed,
    i mod 2      : dense,
    j mod 3      : dense
  )
}>
... tensor<20x30xf32, #BSR> ...

// Same block sparse row storage (2x3 blocks) but this time
// also with a redundant reverse mapping, which can be inferred.
#BSR_explicit = #sparse_tensor.encoding<{
  map = { ib, jb, ii, jj }
        ( i = ib * 2 + ii,
          j = jb * 3 + jj) ->
  ( ib = i floordiv 2 : dense,
    jb = j floordiv 3 : compressed,
    ii = i mod 2 : dense,
    jj = j mod 3 : dense)
}>
... tensor<20x30xf32, #BSR_explicit> ...

// ELL format.
// In the simple format for matrix, one array stores values and another
// array stores column indices. The arrays have the same number of rows
// as the original matrix, but only have as many columns as
// the maximum number of nonzeros on a row of the original matrix.
// There are many variants for ELL such as jagged diagonal scheme.
// To implement ELL, map provides a notion of "counting a
// dimension", where every stored element with the same coordinate
// is mapped to a new slice. For instance, ELL storage of a 2-d
// tensor can be defined with the mapping (i, j) -> (#i, i, j)
// using the notation of [Chou20]. Lacking the # symbol in MLIR's
// affine mapping, we use a free symbol c to define such counting,
// together with a constant that denotes the number of resulting
// slices. For example, the mapping [c](i, j) -> (c * 3 * i, i, j)
// with the level-types ["dense", "dense", "compressed"] denotes ELL
// storage with three jagged diagonals that count the dimension i.
#ELL = #sparse_tensor.encoding<{
  map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
}>
... tensor<?x?xf64, #ELL> ...

// CSR slice (offset = 0, size = 4, stride = 1 on the first dimension;
// offset = 0, size = 8, and a dynamic stride on the second dimension).
#CSR_SLICE = #sparse_tensor.encoding<{
  map = (i : #sparse_tensor<slice(0, 4, 1)>,
          j : #sparse_tensor<slice(0, 8, ?)>) ->
        (i : dense, j : compressed)
}>
... tensor<?x?xf64, #CSR_SLICE> ...