MLIR의 희소 텐서 유형

희소 텐서 인코딩은 정보를 인코딩하는 텐서의 희소성 속성 기반 희소 텐서의 TACO 공식화로 변환할 수 있습니다. 이 인코딩은 결국 Sparsifier 패스를 사용하여 희소 코드를 생성하며, 희소성에 구애받지 않는 계산 표현(즉, 암시적 희소 데이터) 표현은 다음과 같은 명시적 희소 표현으로 변환됩니다. 동시 반복 루프는 희소 스토리지에서 작동하는 희소성이 있는 텐서 대신 인코딩합니다. 이 스파러파이어 패스 전에 실행되는 컴파일러 패스는 인식해야 합니다. 텐서 유형의 의미 체계를 정의합니다.

이 인코딩에서는 크기를 사용하여 시맨틱 텐서의 축을 나타내며 및 level은 실제 저장 형식의 축을 나타냅니다(예: 희소 텐서의 작업 표현을 생성합니다. 개수 차원은 일반적으로 레벨 수 (예: CSR 스토리지 형식) 그러나 인코딩은 더 높은 순서 수준 (예: 블록 희소 BSR 저장 형식을 인코딩하거나)보다 낮은 순서의 예를 들어 스토리지의 단일 레벨로 측정기준을 선형화합니다.

인코딩에는 다음을 제공하는 지도가 포함됩니다.

  • 측정기준 사양의 순서가 지정된 순서로, 각각은 다음을 정의합니다. <ph type="x-smartling-placeholder">
      </ph>
    • 차원 크기 (텐서의 차원 형태에서 암시적)
    • 측정기준-표현식
  • 레벨 사양의 순서가 지정된 시퀀스이며, 각 시퀀스에는 필수 level-type: 레벨 저장 방법을 정의합니다. 각 레벨 유형은 다음으로 구성됩니다. <ph type="x-smartling-placeholder">
      </ph>
    • 저장되는 내용을 정의하는 수준-표현식
    • level-format
    • level-format에 적용되는 level-properties 컬렉션

각 수준-표현식은 아핀 표현식입니다. 사용하는 것이 좋습니다 따라서 수준-표현식은 아핀 맵은 차원 좌표에서 레벨 좌표입니다. 측정기준-표현식은 역지도를 총체적으로 정의하고 추론할 수 없는 복잡한 사례에만 제공하면 됩니다. 자동으로 확장 및 축소할 수 있습니다

각 측정기준에 선택사항인 SparseTensorDimSliceAttr가 있을 수도 있습니다. 희소 저장 형식 내에서 명시적으로 저장된 색인은 좌표 및 오프셋을 positions로 저장 형식으로 반환합니다.

지원되는 수준 형식은 다음과 같습니다.

  • dense : 이 레벨의 모든 항목이 저장됩니다.
  • 압축 : 이 수준에서 0이 아닌 값만 저장됩니다.
  • loose_compressed : 압축된 상태로, 리전 간 여유 공간이 허용되므로
  • singleton : 압축된 형식의 변형으로, 좌표에 동위 요소가 없습니다.
  • block2_4 : 압축 시 1x4 블록당 2:4 인코딩을 사용합니다.

압축된 수준의 경우 각 위치 간격은 축약된 형태로 표시됩니다. 하한 pos(i)과 상한 pos(i+1) - 1이 있는 방식으로, 이는 다음을 의미합니다. 연속적인 간격은 '구멍' 없이 순서대로 나타나야 합니다. 사이 있습니다. 느슨한 압축 형식은 각 요소를 표현하여 이러한 제약 조건을 완화합니다. 하한 lo(i) 및 상한 hi(i)로 된 위치 간격입니다. 간격을 임의의 순서로 표시하고 팔꿈치 사이에 간격을 둘 수 있습니다.

기본적으로 각 레벨 유형은 고유한 속성을 갖습니다 (중복되지 않음). 해당 수준의 좌표) 및 순서가 지정된 (좌표가 해당 수준에서 정렬된 상태로 표시됨) 있습니다. 다음 속성을 수준 형식에 추가하여 변경할 수 있습니다. 다음과 같습니다.

  • 고유하지 않음 : 레벨에 중복된 좌표가 표시될 수 있습니다.
  • 순서 없음 : 좌표가 수수로 표시될 수 있음

맵 외에도 다음 두 필드는 선택사항입니다.

  • 위치 저장소에 필요한 비트폭 (적분 오프셋) 희소 저장소 스키마로 대체합니다. 너비가 좁으면 메모리가 줄어들고 단, 너비가 1,000바이트보다 더 큰 공간을 차지할 수 있는 한 필요한 전체 범위 (즉, 저장된 모든 간접 수준 엔트리가 포함됨). 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에 매핑됨을 나타냅니다. 이는 크기 4로 표시되는 dense 수준입니다. 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 |     |
 +-----+-----+-----+    +-----+-----+-----+

이것은 궁극적으로

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

부수적으로 말 그대로 cuSparse 문서에 나와 있는 NVidia 블록 형식

블록 희소 행. (n.d.-b). Nvidia 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 문서에 제공된 샘플 매트릭스는 다음과 같습니다.

희소 MMA 저장소 예. (n.d.). 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> ...