MLIR 中的稀疏張量類型

稀疏張量編碼是一種屬性,用於將資訊編碼 這些靈感來自於張量的稀有屬性 除以稀疏張量的 TACO 正規化處理次數這個編碼最終是 比對器傳遞,以便從 沒有通用的算式表示法,亦即隱式稀疏 形式會轉換成明確的稀疏表示法 在稀疏儲存空間上執行共同疊代迴圈 而非張量 編碼。在此備用票證之前執行的編譯器票證需要瞭解 這種較稀疏的編碼方式表示了張量型別的語意。

在此編碼中,我們使用「dimension」 就是語意張量的軸 和 level 代表實際儲存空間格式的軸,也就是 記憶體中稀疏張量的操作表示法。開啟 通常會與 等級數量 (例如 CSR 儲存空間格式)。 不過,編碼也能對應 設為排序較高的層級 (例如 對區塊稀疏 BSR 儲存格式進行編碼,或編碼至較低階等級 (例如,將維度線性化為儲存空間中的單一層)。

編碼包含的地圖,提供以下項目:

  • 有排序的維度規格序列,每個規格都有定義:
    • 尺寸大小 (從張量的維度形狀隱含)
    • 維度運算式
  • 按照順序排列的等級規格,每個等級都包含必要設定 level-type:定義等級的儲存方式。每個層級類型 包含:
    • level-expression,定義要儲存的內容
    • 層級格式
    • 適用於關卡格式的層級屬性集合

每個等階運算式都是一個擬因運算式 勝過維度變數因此, 定義 建立來自維度座標的肯定地圖 等級座標。維度運算式 共同定義反向地圖 而這只需針對無法推斷出的詳盡案例提供 。

每個維度也可以有選用的SparseTensorDimSliceAttr。 在稀疏儲存空間格式中,我們是指明確儲存的索引 做為座標,並稱為儲存格式,做為位置

支援的等級格式如下:

  • dense:儲存這個層級的所有項目
  • compress:只會儲存這個層級的非零
  • loose_compressed :壓縮、 但區域之間的空間
  • 單例模式:壓縮格式的變化版本,其中座標沒有同層級
  • block2_4:壓縮會使用每個 1x4 區塊的 2:4 編碼

針對壓縮層級,每個位置間隔都會以精簡格式表示 範圍下限 pos(i) 和上限 pos(i+1) - 1 以表示 連續間隔必須按順序排列,不含任何「孔」在其中 具體做法是指示 Kubernetes 建立並維護 一或多個代表這些 Pod 的物件寬鬆的壓縮格式會代表各項限制,以放寬這些限制 位置間隔,下限 lo(i) 和上限 hi(i), 可讓間隔以任意順序顯示,中間長有肘部空間。

根據預設,每個層級類型的屬性 (不可重複) 以該等級的座標) 及排序 (其座標顯示 層級)。你可以在等級格式中加入下列屬性來變更 以下預設的行為:

  • 非不重複:可能會在樓層顯示重複的座標
  • notordered:座標可能會顯示為任意排列順序

除了地圖,下列兩個是選填欄位:

  • 位置儲存空間所需的位元寬度 (內建偏移量) 放入稀疏儲存配置)。窄寬度會減少記憶體用量 佔用的儲存空間 定義所需的總範圍 (viz. 儲存檔案的數量上限) 項目。選項包括 8163264 或預設的 0 表示原生位元寬度。

  • 座標儲存空間所需的位元寬度 (座標 )。窄寬度會減少記憶體用量 因此即使寬度夠長 總共需要的範圍 (每個張量的最大值 viz. )。選項包括 8163264 或預設的 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

基本上, cuSparse 說明文件中呈現的 NVidia 區塊格式。

封鎖稀疏資料列。 (不適用日期)。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> ...