תרגום שיחה עם StableHLO

המטרה העיקרית של המתרגם של StableHLO היא לספק לסמנטיקה של StableHLO, בהתאם למפרט. היעד המשני הוא שההטמעה תעקוב מקרוב המפרט, עם עדיפות לקריאות על פני ביצועים, כדי לספק בהירות נוספת לסמנטיקה של אפילו של הפעולות המעורבות ביותר כמו Convolution, Gather/Scatter ו-DotGeneral.

נכון לעכשיו, OpenXLA תומך בפרשנות של 91 מתוך 96 מוגדרות מראש פעולות StableHLO. ב-3 הפעולות שנותרו (FftOp, RngOp, RngBitGeneratorOp) הסמנטיקה שלהם מתועדת spec.md, בדיקות ראשוניות שבוצעו במטרה להתקדם (מידע נוסף זמין כאן: status.md כדי לקבל רשימה מלאה של הפעולות והסטטוס העדכני ביותר שלהן). המשחקים הסופיים האלה השיפורים ייושמו על בסיס קהילה לפי הצורך.

היקף

סיווגנו את פעולת StableHLO ל-11 קטגוריות שכוללות 118 פעולות בסך הכול (ראו נספח). הטמעה של קובצי עזר בארגון העבודה על הטמעת תרגום שיחה עבור 100% מפעולות StableHLO כפי שמוגדר במפרט StableHLO. אנחנו אתם מתכננים להשלים את כל העבודה או את רובה בתהליך העבודה הזה ב-StableHLO. גרסה 1.0. מתוך 96 הפעולות שיש להן מפרט כרגע, אנחנו יכולים לפרש 91 פעולות OpenXLA (לעיון במקרים מיוחדים לגבי 5 הנותרים).

מפרט

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

מקרים מיוחדים

שונות

לקטגוריה הזו יש אפשרויות מתפרקות שהעתיד שלהן לא ברור כרגע. יש יש שלוש פעולות מפוקפקות בקטגוריה הזו שבהן המתרגם לא תומך בהן הרגע:

• FftOp
• RngOp
• RngBitGeneratorOp

FftOp מסווג כ'שונות', אבל בניגוד לתפעול אחר בקטגוריה הזו, לתפעול הזה אין אישור מרחיב, והתמיכה בו ב-StableHLO היא WIP.

אפשר לפרק את RngOp ואת RngBitGeneratorOp לתפעול MHLO, אבל של פירוק כימי XlaRngGetAndUpdateStateOp, שהוא מודל MHLO ספציפי מערכת הפעלה תמיכה בפרשנות של שתי הפעולות האלה היא WIP.

הכלי להמרת הפעולות שנותרו בקטגוריה הזו לתפעול StableHLO המתרגם תומך ב-hlo_expand_main.cc.

לא ב-HLO

מלבד הפעולות שנבחרו, הקטגוריה הזו כוללת 8 פעולות שלא נבדקו (ראו StableHLO Ops categories) שמתוכננות להיות הועבר מ-StableHLO. לרוב הפעולות האלה יש אישורים קיימים mhlo להמיר אותן לפעולות שוות ערך ב-StableHLO.

הכלי להמרת הפעולות שנותרו בקטגוריה הזו לפעולות מקבילות של StableHLO שרכיב התרגום תומך בהן נמצא בכתובת mlir-hlo-opt.cc.

כימות

תמיכת תרגום לפעולה של stablehlo.constant עם סוג כמותי היא אין תמיכה, ובמעקב דרך #1691.

הוראות שימוש

יצירת 'כלי התרגום'

אפשר ליצור את מפענח התרגום ולבדוק אותו באמצעות Bazel או CMake (מועדף). לגרסה מלאה הוראות להתאמה אישית, ראה README.md.

בזל:

bazel build //...

CMake:

mkdir -p build && cd build

cmake .. -GNinja \
  -DLLVM_ENABLE_LLD="$LLVM_ENABLE_LLD" \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=On \
  -DMLIR_DIR=${PWD}/../llvm-build/lib/cmake/mlir

כדי להפעיל את המתורגמן, יש לנו כלי תרגום לתרגום התוכניות StableHLO שכתוב ב-MLIR.

stablehlo-translate --interpret <path/to/program>

ניב התרגום

הניב Interpreter מכיל פעולות שירות שונות שקשורות בתרגום השיחה הפעילה. באופן ספציפי, interpreter.run_parallel (יש לעיין InterpreterOps.td לסמנטיקה של תפעול ושימוש לדוגמה) הוא מאפשר פירוש של פעולות הפצה ועוד יתווספו לתוכנית בהתאם לצרכים של הקהילה.

ניב הבדיקה

הדיאלקט Check משמש להשוואה בין ערכי זמן הריצה של המתרגם לבין הערכים הצפויים ערכים. ניתן לבדוק את הפלטים של תוכנית StableHLO באמצעות פעולות בדיקה שונות (ראו CheckOps.td סמנטיקה של תפעול ושימוש לדוגמה).

תוכניות למבחני כתיבה

אנחנו משתמשים בכלי lit של LLVM כדי להריץ להשוות לקובץ שנוצר כדי להבדיל מהפלט של המתרגם (מידע נוסף זמין בכתובת stablehlo/tests/interpret למשל בדיקות).

מתבצעת בדיקה של AddOp (דוגמה של interpret_add.mlir):

// RUN: stablehlo-translate --interpret %s

func.func @add_op_scalar() {
  %0 = stablehlo.constant dense<2> : tensor<i4>
  %1 = stablehlo.constant dense<3> : tensor<i4>
  %2 = stablehlo.add %0, %1 : tensor<i4>
  check.expect_eq_const %2, dense<5> : tensor<i4>
  func.return
}

כדי לבדוק פעולות בקטגוריית ההפצה, צריך להפעיל אותן דרך הפעלת שירות interpreter.run_parallel

מתבצעת בדיקה של AllReduceOp (דוגמה של all_reduce.mlir):

// RUN: stablehlo-translate --interpret %s

module @cross_replica {
  func.func public @all_reduce(%operand : tensor<4xi64>) -> tensor<4xi64> {
    %result = "stablehlo.all_reduce"(%operand) ({
      ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
        %0 = stablehlo.add %arg0, %arg1 : tensor<i64>
        stablehlo.return %0 : tensor<i64>
    }) {
      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
      channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
    } : (tensor<4xi64>) -> tensor<4xi64>
    return %result : tensor<4xi64>
  }
  func.func public @main() {
    %inputs0 = stablehlo.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %inputs1 = stablehlo.constant dense<[5, 6, 7, 8]> : tensor<4xi64>
    %results:2 = "interpreter.run_parallel"(%inputs0, %inputs1) {
      programs=[[@all_reduce], [@all_reduce]]
    } : (tensor<4xi64>, tensor<4xi64>) -> (tensor<4xi64>, tensor<4xi64>)
    check.expect_eq_const %results#0, dense<[6, 8, 10, 12]> : tensor<4xi64>
    check.expect_eq_const %results#1, dense<[6, 8, 10, 12]> : tensor<4xi64>
    func.return
  }
}

ניפוי באגים ב-StableHLO

בהתאם לשלבי ה-build של StableHLO, הקבצים הבינאריים של StableHLO לכלים stablehlo/tools צריך להתגורר ב-/build/bin. כלים נפוצים לניפוי באגים, כמו אפשר להשתמש ב-GDB כדי לעבור את הקוד:

gdb --args ./build/bin/stablehlo-translate -allow-unregistered-dialect --interpret ./stablehlo/tests/interpret/<test>.mlir

נספח

המרת פעולות שונות

# batch_norm_grad
hlo-expand --batch_norm_grad_expander <path/to/hlo_module>

# batch_norm_inference
hlo-expand --batch_norm_inference_expander <path/to/hlo_module>

# batch_norm_training
hlo-expand --batch_norm_training_expander <path/to/hlo_module>

# cholesky
hlo-expand --cholesky_expander <path/to/hlo_module>

# constant
# Supported in StableHLO interpreter.

# fft
# TBD

# iota
# Supported in StableHLO interpreter.

# rng
# TBD

# rng_bit_generator
# TBD

# triangular_solve
hlo-expand --triangular_solve_expander <path/to/hlo_module>

המרה לפעולות 'לא ב-HLO'

# broadcast
mlir-hlo-opt -mhlo-legalize-broadcast-to-broadcast-in-dim <path/to/input>

# create_token
mlir-hlo-opt -mhlo-legalize-create-token-to-after-all <path/to/input>

# cross-replica-sum
mlir-hlo-opt -mhlo-legalize-cross-replica-sum-to-all-reduce <path/to/input>

# dot
mlir-hlo-opt -mhlo-legalize-dot-to-dot-general <path/to/input>

# einsum
mlir-hlo-opt -mhlo-legalize-einsum-to-dot-general <path/to/input>

# torch_index_select
mlir-hlo-opt -mhlo-legalize-torch-index-select-to-gather <path/to/input>

# unary_einsum
mlir-hlo-opt --canonicalize -mhlo-legalize-einsum-to-dot-general <path/to/input>

קטגוריות תפעול StableHLO