此页面由 Cloud Translation API 翻译。

MLIR 中的稀疏张量类型

稀疏张量编码是用于对信息进行编码的属性关于张量的稀疏属性，灵感来自于因为稀疏张量的 TACO 规范化而受到影响。这种编码最终会用到 sparsifier 传递，以通过计算的稀疏性表示法，即隐式稀疏被转换为显式稀疏表示法，其中协同迭代循环在稀疏存储上运行而不是具有稀疏性的张量，编码。需要了解在此稀疏器传递之前运行的编译器传递以及这种稀疏性编码的张量类型语义。

在此编码中，我们使用 dimension 来指语义张量的轴，和 level 共同表示实际存储格式的轴，即表示内存中的稀疏张量的运算表示法。您获得的维度通常与级别数（例如 CSR 存储格式）。不过，编码还可以将高阶维度（例如，对块稀疏 BSR 存储格式进行编码）或编码为低阶级别的 BSR （例如，将维度线性化为存储空间中的单个级别）。

编码包含的映射提供以下内容：

一系列有序的尺寸规范，其中每个规范都定义： <ph type="x-smartling-placeholder">
- 维度大小（隐含自张量的维度形状）
- 维度表达式
一系列有序的等级规范，其中每个等级规范都包含一个 level-type：用于定义关卡的存储方式。每种关卡类型包括： <ph type="x-smartling-placeholder">
- 一个 level-expression，定义存储的内容
- 级别格式
- 一组适用于 level-format 的 level-properties

每个级别表达式都是一个仿射表达式对维度变量进行过滤。因此， level-expressions 共同定义从维度坐标到级别坐标。维度表达式共同定义反向地图，只有在无法得出结论的情况。

每个维度还可以有一个可选的 SparseTensorDimSliceAttr。在稀疏存储格式中，我们是指显式存储的索引表示为坐标，偏移量表示为 positions。

支持的级别格式如下：

dense：存储此级别的所有条目
compressed：仅存储此级别的非零值
loose_compressed：压缩格式但每个区域之间
单例：压缩格式的变体，其中坐标没有同级元素
block2_4：对于每个 1x4 块，压缩使用 2:4 编码

对于压缩级别，每个位置间隔都以紧凑的形式表示，设置下限 pos(i) 和上限 pos(i+1) - 1，这意味着连续间隔必须按顺序出现，且不能出现任何“空洞”介于。松散压缩格式通过表示排名范围，其中下限为 lo(i)，上限为 hi(i)，让间隔可按任意顺序显示，且之间可留出肘部空间。

默认情况下，每个关卡类型都具有唯一性（不重复）坐标）和排序（坐标显示时按级别）。将以下属性添加到级别格式可更改这一默认行为：

nonunique：级别上可能会出现重复的坐标
无序：坐标可按任意顺序出现

除了映射之外，还有以下两个字段是可选字段：

位置存储所需的位宽（整数偏移）转换为稀疏存储方案）。较窄的宽度会减少内存占用的存储空间，只要宽度足以定义所需的总范围（即所存储的条目）。选项包括 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（以 position 数组和坐标数组。值 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

顺便说一句 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 存储示例。（注释）。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> ...

MLIR 中的稀疏张量类型 使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

示例

MLIR 中的稀疏张量类型