Tipos de tensores esparsos no MLIR

A codificação de tensor esparso é um atributo para codificar informações nas propriedades de esparsidade dos tensores, inspiradas pela formalização do TACO de tensores esparsos. Essa codificação é usada por uma passagem sparsifier para gerar código esparso de maneira totalmente automática a partir de uma representação independente de esparsidade do cálculo, ou seja, uma representação esparsa implícita é convertida em uma representação esparsa explícita em que os loops de coiteração operam no armazenamento esparso em vez de tensores com esparsidade e codificação. Os passes do compilador executados antes dessa passagem do esparsador precisam estar cientes da semântica de tipos de tensores com essa codificação de esparsidade.

Nesta codificação, usamos dimension para consulte os eixos do tensor semântico, e level para se referir aos eixos do formato de armazenamento real, ou seja, o representação operacional do tensor esparso na memória. O número de normalmente é igual a o número de níveis (como o formato de armazenamento CSR). No entanto, a codificação também pode mapear para níveis de ordem superior (por exemplo, para codificar um formato de armazenamento BSR esparso em blocos) ou para níveis de ordem inferior Por exemplo, para linearizar as dimensões como um único nível no armazenamento.

A codificação contém um mapa que fornece o seguinte:

  • Uma sequência ordenada de especificações de dimensão, cada uma definindo:
    • o tamanho da dimensão (implícito do formato de dimensão do tensor)
    • uma dimension-expression;
  • Uma sequência ordenada de especificações de nível, cada uma incluindo uma função level-type, que define como armazenar o nível. Cada tipo de nível consiste em:
    • uma level-expression, que define o que é armazenado
    • um formato de nível
    • um conjunto de propriedades de nível que se aplicam ao formato de nível

Cada expressão de nível é afins sobre as variáveis de dimensão. Assim, o as expressões de nível definem coletivamente mapa afinado das coordenadas de dimensão para coordenadas de nível. As dimension-expressions definem coletivamente o mapa inverso, que só precisa ser fornecido para casos elaborados em que não é possível inferir automaticamente.

Cada dimensão também pode ter um SparseTensorDimSliceAttr opcional. Dentro do formato de armazenamento esparso, nos referimos a índices armazenados explicitamente como coordenadas e deslocamentos no formato de armazenamento como positions.

Os formatos de nível compatíveis são os seguintes:

  • dense : todas as entradas nesse nível são armazenadas
  • compactado : somente valores diferentes de zero ao longo desse nível são armazenados;
  • loose_compressed : conforme compactado, mas permite espaço livre entre regiões
  • singleton : uma variante do formato compactado, em que as coordenadas não têm irmãos
  • block2_4 : a compactação usa uma codificação 2:4 por bloco 1x4.

Para um nível compactado, cada intervalo de posição é representado em um formato com um limite inferior pos(i) e um pos(i+1) - 1 superior, o que implica que intervalos sucessivos precisam aparecer em ordem, sem nenhum "buraco" entre para resolvê-los com rapidez. O formato compactado solto relaxa essas restrições, representando cada intervalo de posição com um limite inferior lo(i) e um limite superior hi(i), que permite que intervalos apareçam em ordem arbitrária e com espaço entre eles.

Por padrão, cada tipo de nível tem a propriedade de ser único (sem duplicatas coordenadas naquele nível) e ordenadas (as coordenadas aparecem classificadas naquele nível nível de conta). As propriedades a seguir podem ser adicionadas a um formato de nível para mudar esse comportamento padrão:

  • nonunique : é possível que coordenadas duplicadas apareçam no nível
  • nonordered : as coordenadas podem aparecer em ordem arbribratry

Além do mapa, os dois campos a seguir são opcionais:

  • A largura de bits necessária para armazenamento de posições (deslocamentos integrais) para o esquema de armazenamento esparso). Uma largura estreita reduz a memória pegada de armazenamento suspenso, desde que a largura seja suficiente define o intervalo total exigido (visto o número máximo de objetos entradas em todos os níveis de indireção). As opções são 8, 16, 32, 64 ou, o padrão, 0 para indicar a largura de bits nativa.

  • A largura de bits necessária para o armazenamento de coordenadas (as coordenadas de entradas armazenadas). Uma largura estreita reduz o consumo de memória de armazenamento suspenso, contanto que a largura seja suficiente para definir o intervalo total necessário (viz. o valor máximo de cada tensor coordenadas em todos os níveis). As opções são 8, 16, 32, 64 ou, o padrão, 0 para indicar uma largura de bits nativa.

Exemplos

No formato CSR(Compressed Sparse Row) mostrada abaixo, a codificação de tensor esparso seria:

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

Indica que a first dimension (linha) é mapeada para o first level, que é um nível de dense, indicado pelo tamanho 4. E a solicitação second dimension (coluna) mapeia para o second level, indicado pela matriz de posições e de coordenadas. O valor 3 ([1, 1] na matriz original) é representado pelo deslocamento da matriz de posições (valor 3 é o número da linha é 1 na matriz original, pois é o segundo par de deslocamento e seu o número da coluna fica no índice [2 : 4) na matriz de coordenadas). E, no de coordenadas, podemos ver que o número da coluna 3 do valor é 1 na matriz original.

Para o formato BSR(Block Sparse Row), o tipo de tensor esparso é:

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

Considere a seguinte matriz esparsa, com blocos 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 |     |
 +-----+-----+-----+    +-----+-----+-----+

que é armazenado no formato com sabor do TACO como

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

Que é literalmente a Formato de bloco NVidia apresentado na documentação cuSparse.

Bloquear linha esparsa. (sem data - b). NVIDIA https://docs.nvidia.com/cuda/cusparse/_images/bsr.png

Também aceitamos Nvidia's 2:4 structured sparsity.

O tipo de tensor esparso para isso é o seguinte:

#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
}>

Dada a matriz de amostra da documentação da NVidia:

Exemplo de armazenamento de MMA esparsa. (sem data). NVIDIA https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-sparse-mma

O MLIR agora mapeia essa matriz para um layout idêntico:

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

Mais exemplos:

// 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> ...