דינמיות ב-StableHLO

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

סקירה כללית של מונחי דינאמיזם ותמיכה

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

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

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

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

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

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

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

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

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

פולימורפיזם לצורה

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

יש שתי השלכות עיקריות על יצירת פולימורפיזם:

1. כל הדינמיות בתוכנית מתייחסת לפרמטרים של הקלט שלה.
2. כל הדינמיות מתייחסת לצורות של הטנזורים בלבד, כלומר היא לא תלויה בנתונים.

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

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

דינמיות שמבוססת על נתונים

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

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

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

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

• מדריך JAX > ייצוא עם גודל קבוצה דינמי
• מדריך PyTorch/XLA > ייצוא באצווה דינמית

שלבי מהדר לשיפור תוכניות דינמיות

הסרת צינור עיבוד הנתונים של Dynamism Pass

יש כמה שלבים שימושיים לשיפור הצורות, והם כולם נמצאים בצינור עיבוד נתונים של שלבים createStablehloRemoveDynamismPipeline:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

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

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

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

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

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

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

מודל 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>. אם נשנה את גודל האצווה או את אורך הרצף, נצטרך לעקוב מחדש אחרי קוד המקור ולהוריד אותו מחדש ל-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>
}