Разреженные типы тензоров в MLIR

Разреженное тензорное кодирование — это атрибут для кодирования информации о свойствах разреженности тензоров, вдохновленный формализацией разреженных тензоров TACO. Это кодирование в конечном итоге используется проходом разрежителя для полной автоматической генерации разреженного кода из независимого от разреженности представления вычислений, т. е. неявное разреженное представление преобразуется в явное разреженное представление, где совместные циклы работают с разреженными форматами хранения, а не с разреженными форматами хранения. тензоры с разреженной кодировкой. Проходы компилятора, которые выполняются перед этим проходом разрежителя, должны учитывать семантику тензорных типов с такой разреженной кодировкой.

В этом кодировании мы используем размерность для обозначения осей семантического тензора, а уровень — для обозначения осей фактического формата хранения, т. е. рабочего представления разреженного тензора в памяти. Количество измерений обычно совпадает с количеством уровней (например, формат хранения CSR). Однако кодирование также может отображать измерения на уровни более высокого порядка (например, для кодирования формата хранения BSR с разрежением блоков) или на уровни более низкого порядка (например, для линеаризации измерений как одного уровня в хранилище).

Кодировка содержит карту, которая обеспечивает следующее:

  • Упорядоченная последовательность спецификаций размеров, каждая из которых определяет:
    • размер-размер (неявно вытекающий из размерной формы тензора)
    • выражение-размерность
  • Упорядоченная последовательность спецификаций уровней, каждая из которых включает требуемый тип уровня , определяющий, как уровень должен храниться. Каждый тип уровня состоит из:
    • выражение уровня , которое определяет, что хранится
    • формат уровня
    • коллекция свойств уровня , которые применяются к формату уровня

Каждое выражение уровня является аффинным выражением для переменных измерения. Таким образом, выражения уровня коллективно определяют аффинную карту от координат измерения до координат уровня. Выражения размерностей в совокупности определяют обратную карту, которую необходимо предоставлять только для сложных случаев, когда ее невозможно вывести автоматически.

Каждое измерение также может иметь необязательный SparseTensorDimSliceAttr . В разреженном формате хранения мы ссылаемся на индексы, которые хранятся явно как координаты , а смещения в формате хранения — как позиции .

Поддерживаются следующие форматы уровней:

  • плотный : все записи на этом уровне сохраняются.
  • сжатый : на этом уровне сохраняются только ненулевые значения.
  • «loose_compressed» : как сжатый, но допускает свободное пространство между регионами.
  • Singleton : вариант сжатого формата, в котором координаты не имеют братьев и сестер.
  • block2_4 : при сжатии используется кодировка 2:4 для каждого блока 1x4.

Для сжатого уровня каждый интервал позиции представлен компактно с нижней границей pos(i) и верхней границей pos(i+1) - 1 , что означает, что последовательные интервалы должны появляться по порядку без каких-либо «дыр» между ними. . Свободный сжатый формат ослабляет эти ограничения, представляя каждый интервал позиции с нижней границей lo(i) и верхней границей hi(i) , что позволяет интервалам появляться в произвольном порядке и с небольшим пространством между ними.

По умолчанию каждый тип уровня имеет свойство быть уникальным (на этом уровне нет повторяющихся координат) и упорядоченным (координаты отображаются отсортированными на этом уровне). Следующие свойства можно добавить в формат уровня, чтобы изменить поведение по умолчанию:

  • неуникальный : на уровне могут появиться повторяющиеся координаты.
  • неупорядоченный : координаты могут отображаться в произвольном порядке.

Помимо карты, следующие два поля являются необязательными:

  • Требуемая разрядность для хранения позиций (интегральные смещения в разреженной схеме хранения). Узкая ширина уменьшает объем памяти, занимаемой служебным хранилищем, при условии, что ширины достаточно для определения общего требуемого диапазона (т. е. максимального количества сохраненных записей на всех уровнях косвенности). Возможные значения: 8 , 16 , 32 , 64 или, по умолчанию, 0 , чтобы указать собственную разрядность.

  • Требуемая разрядность для хранения координат (координат сохраняемых записей). Узкая ширина уменьшает объем памяти, занимаемой служебным хранилищем, при условии, что ширины достаточно для определения общего требуемого диапазона (т. е. максимального значения каждой координаты тензора на всех уровнях). Возможные значения: 8 , 16 , 32 , 64 или, по умолчанию, 0 , чтобы указать собственную разрядность.

Примеры

В формате CSR(Compressed Sparse Row) показанном ниже, разреженная тензорная кодировка будет такой:

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

Это указывает на то, что first dimension (строка) сопоставляется с first level , который представляет собой dense уровень, обозначенный размером 4. А second dimension (столбец) сопоставляется со second level , обозначенным массивом позиций и массивом координат. Значение 3 ( [1, 1] в исходной матрице) представлено смещением от массива позиций (номер строки значения 3 равен 1 в исходной матрице, поскольку это вторая пара смещений, а номер ее столбца находится в индексе [2 : 4) в массиве координат). И в массиве координат мы видим, что номер столбца значения 3 равен 1 в исходной матрице.

Для формата BSR(Block Sparse Row) тип разреженного тензора:

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

Рассмотрим следующую разреженную матрицу с блоками 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 |     |
 +-----+-----+-----+    +-----+-----+-----+

который в конечном итоге сохраняется в формате TACO как

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

Это, кстати, буквально блочный формат NVidia, представленный в документации cuSparse.

Блокировать разреженную строку . (нд-б). Нвидиа. https://docs.nvidia.com/cuda/cusparse/_images/bsr.png

Мы также поддерживаем Nvidia's 2:4 structured sparsity .

Разреженный тип тензора для этого выглядит следующим образом:

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

Учитывая пример матрицы, приведенный в документации NVidia:

Пример разреженного хранилища ММА . (н-й). Нвидиа. https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-for-sparse-mma

MLIR теперь сопоставляет эту матрицу с идентичным макетом:

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

Еще примеры:

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