שיטת Tensor דחופה ב-MLIR

קידוד טנזור קטן הוא מאפיין לקידוד מידע על תכונות מיעוט של טנזור, על ידי צורת הפורמליזציה של TACO של מקטעים דלים. בסופו של דבר נעשה שימוש בקידוד הזה באמצעות מעבר sparsifier כדי ליצור קוד קטן לחלוטין באופן אוטומטי ייצוג שלא תלוי-נתונים של החישוב, כלומר נתח קטן מרומז שהייצוג שלו מומר לייצוג מועט, שבו לולאות חזרה משותפות פועלות על נפח אחסון קטן במקום tensors עם מעט מאוד באמצעות הקידוד. חשוב לשים לב לכרטיסי מהדר שמריצים לפני אישור המגביל הזה מהסמנטיקה של סוגי הטנזור עם קידוד כזה.

בקידוד הזה אנחנו משתמשים במאפיין כדי מתייחס לצירים של הטנסור הסמנטי, ו-level כדי להתייחס לצירים של פורמט האחסון בפועל, כלומר או ייצוג תפעולי של הטנזור הנדיר בזיכרון. מספר המאפיינים בדרך כלל זהים מספר הרמות (למשל פורמט האחסון של CSR). אבל הקידוד יכול גם למפות לרמות גבוהות יותר (לדוגמה, כדי לקודד פורמט אחסון BSR עם חסימת בלוקים) או לרמות נמוכות יותר של (לדוגמה, כדי לבצע ליניארית מאפיינים כרמה אחת באחסון).

הקידוד מכיל מפה שמספקת את הפרטים הבאים:

  • רצף מסודר של מפרטי מאפיינים, שכל אחד מהם מגדיר:
    • גודל המימד (משתמע מצורת המימד של הטנזור)
    • ביטוי-מאפיין
  • רצף סדור של מפרטי רמות, כאשר כל אחד מהם כולל שלב level-type (סוג הרמה): אופן האחסון של הרמה. כל סוג רמה מורכב מ:
    • ביטוי רמה, שמגדיר מה נשמר
    • פורמט רמה
    • אוסף של מאפייני רמה שחלים על פורמט הרמה

כל ביטוי רמה הוא ביטוי עניין על פני משתני מאפיין. לכן, ביטויי רמה מגדירים באופן קולקטיבי מפה קשורה מקואורדינטות של ממדים עד הקואורדינטות של הרמה. ביטויי המאפיינים מגדירים באופן קולקטיבי את המפה ההפוכה, יש לספק רק במקרים מורכבים שבהם לא ניתן להסיק באופן אוטומטי.

לכל מאפיין יכול להיות גם SparseTensorDimSliceAttr אופציונלי. בפורמט של נפח אחסון קטן, אנחנו מתייחסים למדדים שמאוחסנים במפורש כקואורדינטות, שמקבילות לפורמט האחסון כמיקומים.

הפורמטים הנתמכים של הרמות הם:

  • צפיפות : כל הרשומות ברמה הזו נשמרות
  • דחוס : רק פריטים חד-פעמיים ברמה הזו נשמרים
  • loose_compressed : דחוס, אבל מאפשר מרווח פנוי בין אזורים
  • singleton : וריאציה של הפורמט הדחוס, שבו לקואורדינטות אין אחים
  • block2_4 : הדחיסה משתמשת בקידוד 2:4 לכל בלוק 1x4

ברמה דחוסה, כל מרווח מיקום מיוצג בעמודה קומפקטית. דרך עם pos(i) גבול תחתון ו-pos(i+1) - 1 גבול עליון, שמרמז על גבול תחתון שמרווחי זמן עוקבים חייבים להופיע בסדר ללא "חורים" בין אותם. הפורמט הדחוס החופשי משחרר את המגבלות האלה על ידי ייצוג של כל מרווח מיקום עם גבול תחתון lo(i) ו-hi(i) גבול עליון, אשר מאפשרת הצגה של מרווחים לפי סדר שרירותי, עם מרווח מרפק ביניהם.

כברירת מחדל, לכל סוג רמה יש את המאפיין שהוא ייחודי (ללא כפילויות את הקואורדינטות ברמה הזו) ומסודרות (הקואורדינטות מופיעות ממוינות לפי רמה). אפשר להוסיף את המאפיינים הבאים לפורמט של רמה כדי לשנות אותם התנהגות ברירת המחדל הזו:

  • nonUnique : קואורדינטות כפולות עשויות להופיע ברמה
  • nonorder : הקואורדינטות עשויות להופיע בסדר ספרייה

בנוסף למפה, שני השדות הבאים הם אופציונליים:

  • רוחב הסיביות הנדרש לאחסון מיקום (קיזוזים בלתי נפרדים) לסכימת האחסון המצומצמת). רוחב צר מפחית את כמות הזיכרון טביעת הרגל הפחמנית של האחסון העליון, כל עוד הרוחב מספיק להגדיר את הטווח הכולל הנדרש (כלומר את המספר המקסימלי של יחידות מודעות ששמורות רשומות בכל רמות העקיפה). האפשרויות הן 8, 16, 32, 64 או ברירת המחדל, 0, לציון רוחב הביטים המקורי.

  • רוחב הסיביות הנדרש לאחסון קואורדינטות (הקואורדינטות מתוך הרשומות השמורות). רוחב צר מצמצם את טביעת הרגל הפחמנית של שטח האחסון התקורה, כל עוד הרוחב מספיק כדי להגדיר הטווח הנדרש הכולל (כלומר הערך המקסימלי של כל טנזור). וכל הרמות). האפשרויות הן 8, 16, 32, 64, או ברירת המחדל, 0 כדי לציין רוחב סיביות מותאם.

דוגמאות

בפורמט CSR(Compressed Sparse Row) שמוצגת בהמשך, הקידוד של החלק הנדיר (stensor) תהיה:

#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), הסוג של המפריד sparse הוא:

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

חסימת השורה המצומצמת. (י.ד.-ב). 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> ...