דינמיות ב-StableHLO

המצב הנוכחי של הדינמיות מפורט יותר בRFC בנושא דינמיות. בדף הזה מופיע סקירה כללית של ה-RFC, ומוסבר על ממשקי API חשובים וכלים ליצירת אינטראקציה עם תוכניות דינמיות.

מינוח דינמי וסקירה כללית על תמיכה

קודם כל, נסביר כמה מונחים שיופיעו במסמך הזה, וגם ניתן מבוא קצר לתמיכה שלהם ב-StableHLO:

מאפיינים דינמיים

מאפיינים דינמיים הם מאפיינים שגודל המאפיין שלהם לא ידוע. ב-StableHLO, אנחנו מייצגים מאפיינים דינמיים באמצעות ?, כלומר tensor<16x?xf32>.

דינמיות מוגבלת

דינמיות מוגבלת מתייחסת למאפיין דינמי שהערך שלו מוגבל מלמעלה. בדרך כלל, הפעולה הזו שימושית להוספת ריפוד לטנזור במהלך ההפעלה. ב-StableHLO אנחנו מייצגים דינמיות מוגבלת באמצעות #stablehlo.bounds כקידוד טנסור, כלומר טנסור מדרגה 2 עם מימד דינמי אחד שמוגבל ל-16 והשני ללא הגבלה, שאפשר לייצג אותו כ-tensor<?x?xf32, #stablehlo.bounds<16, ?>>.

‫StableHLO יכול לייצג דינמיות מוגבלת, אבל יש תמיכה מוגבלת ב-framework, שמקורה ב-TensorFlow, וקיימת תמיכה מסוימת ב-PyTorch/XLA.

דינמיות בלתי מוגבלת

דינמיות בלתי מוגבלת, כפי שהשם מרמז, מתייחסת למאפיין דינמי ללא גבול ידוע לגודל. הדינמיות הזו נפוצה מאוד ב-StableHLO, עם תמיכה ב-JAX, ב-PyTorch/XLA וב-TF, ומשמשת לעיתים קרובות לייצוא מודלים עם גודל אצווה או אורך רצף דינמיים.

ב-StableHLO אנחנו פשוט משמיטים את קידוד הגבולות עבור הדינמיות מהסוג הזה, כלומר: tensor<?x?xf32>.

פולימורפיזם של צורות

פולימורפיזם של צורה הוא מונח שירשנו מ-JAX.

יש שתי השלכות חשובות לפולימורפיזם של צורות:

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

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

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

דינמיות שתלויה בנתונים

דינמיות שתלויה בנתונים מתייחסת לגדלים דינמיים של מאפיינים שקשורים לנתונים בתוך טנסור. הדוגמה הקנונית היא פונקציית nonzeros שמחזירה את האינדקסים של כל הרכיבים שהם 0 בערך טנסור. אי אפשר לדעת את הצורה בלי להעריך את הנתונים, אבל אפשר לרוב לקמפל אותה באמצעות דינמיות מוגבלת, ולהקצות זיכרון נוסף לגודל הטנסור של הפלט הפוטנציאלי.

אפשר ליצור מודלים של הרבה פעולות דינמיות שתלויות בנתונים באמצעות דינמיות מוגבלת, שבה מציינים חסם עליון לגודל של טנסור, ובדרך כלל החומרה תיישם את זה באמצעות ריפוד טנסור. יש היום תמיכה מסוימת בדינמיות שתלויה בנתונים ב-PyTorch/XLA וב-TensorFlow, אבל JAX לא עוקב כרגע אחרי פעולות שמובילות לדינמיות שתלויה בנתונים.

ייצוא תוכניות עם מאפיינים דינמיים

במדריכים שלנו בנושא StableHLO מוסבר איך לייצא תוכניות עם גדלים דינמיים של אצווה או אורכים דינמיים של רצף:

מעברים של קומפיילר לשיפור תוכניות דינמיות

הסרה של צינור עיבוד נתונים של כרטיס דינמי

יש כמה כרטיסים שימושיים לשיפור הצורות, וכולם כלולים בצינור כרטיסים createStablehloRemoveDynamismPipeline:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

כרטיסים ספציפיים לשיפור הדינמיות

המעברים שמועילים בדרך כלל לשיפור הצורה הם:

‫stablehlo-refine-arguments כדי להחליף את ארגומנטי הקלט בסוגי טנסור קונקרטיים.
‫stablehlo-refine-shapes כדי להפיץ את המידע על הצורה של ארגומנט הקלט החדש בכל התוכנית.
‫stablehlo-canonicalize-dynamism כדי להחליף פעולות דינמיות בווריאציות סטטיות שלהן.
‫stablehlo-check-shape-assertions כדי לבדוק ולהסיר קריאות מותאמות אישית של טענות לגבי צורות.

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

דוגמה: למה דינמיות שימושית ואיך אפשר להשתמש בה?

יש הרבה שימושים לדינמיות, אבל כאן נתמקד בעיקר בתרחיש הנפוץ של פולימורפיזם של צורות – יצירת ייצוג גמיש של מודל מיוצא, שמשמש בדרך כלל לייצוג של גודל אצווה דינמי או אורך רצף.

מודל סטטי של add_one

נשתמש במודל add_one הפשוט הבא כדי להדגים את זה:

def add_one(x):
  return x + 1

כשעוקבים אחריו באמצעות tensor<4xf32>, מקבלים את תוכנית StableHLO הבאה:

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

המודל הזה יפעל רק עבור ארגומנטים של קלט שיש להם צורה של tensor<4xf32> shape. אם נשנה את גודל האצווה או את אורך הרצף, נצטרך לעקוב מחדש אחר קוד המקור ולהמיר אותו מחדש ל-StableHLO, ואין שום ערובה לכך שתהיה לנו גישה לקוד המקור.

מודל דינמי של add_one

כאן נכנסת לתמונה הדינמיות הפולימורפית של הצורה. במקום זאת, JAX ו-PyTorch/XLA יכולים להפיק את מודל add_one עם IR תקף באופן דינמי, שישדר את הקבוע כדי להתאים לצורת הקלט הדינמית באופן הבא:

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

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

שיפור המודל הדינמי

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

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

כך התוכנית משתנה בהדרגה:

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}

דינמיות ב-StableHLO קל לארגן דפים בעזרת אוספים אפשר לשמור ולסווג תוכן על סמך ההעדפות שלך.