انواع تانسور پراکنده در 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) قرار دارد. [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
    )

ماتریس پراکنده زیر را با بلوک‌های ۲×۲ در نظر بگیرید:

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 ارائه شده است.

سطر پراکنده را مسدود کنید . (nd-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:

مثال Sparse MMA Storage . (دوم). 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> ...