כלי XLA

תהליך הפיתוח של XLA מתמקד בדרך כלל ב-IR של HLO, שמייצג חישוב פונקציונלי מבודד שמועבר לקומפיילר. ‫XLA כולל כמה כלים לשורת הפקודה (שמתוארים בהמשך) שצורכים HLO ומריצים אותו, או מספקים שלב ביניים של קומפילציה. השימוש בכלים כאלה הוא בעל ערך רב למחזור איטרציה מהיר, כי אפשר להמחיש את HLO ויש אפשרות לפרוץ אותו, ולעתים קרובות הדרך המהירה ביותר להבין ולתקן ביצועים או התנהגות של XLA היא לשנות אותו ולהריץ אותו שוב ושוב.compile->modify->run

הדרך הכי פשוטה לקבל את ה-HLO של תוכנית שעוברת קומפילציה באמצעות XLA היא בדרך כלל באמצעות משתנה הסביבה XLA_FLAGS:

$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point

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

‫[`run_hlo_module`] הפעלת מודולים של HLO

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

הכלי run_hlo_module פועל על HLO לפני אופטימיזציה, ובאופן שמוגדר כברירת מחדל הוא כולל קומפילציה, הפעלה והשוואה עם יישום המפרש של ההפניה. לדוגמה, הקריאה הרגילה להרצת קובץ קלט computation.hlo ב-GPU של NVIDIA ולבדיקת התקינות שלו היא:

run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

הרצת כמה מודולים של HLO

הפעלת פונקציה עם כמה מודולים של HLO נתמכת ב-run_hlo_module. כדי להריץ את כל מודולי ה-HLO מספרייה:

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

‫[`multihost_hlo_runner`] הרצת מודולים של HLO עם תמיכה ב-SPMD

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

‫Multihost HLO runner הוא כלי דומה מאוד, אבל הוא תומך ב-SPMD, כולל תקשורת בין מארחים. פרטים נוספים זמינים במאמר בנושא הפעלת HLO בכמה מארחים.

הרצת מודולים מרובים של HLO עם תמיכה ב-SPMD

בדומה ל-run_hlo_module, גם multihost_hlo_runner תומך בהפעלה עם כמה מודולים.

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*

‫[`hlo-opt`] הידור מודול HLO

bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>

כשמבצעים ניפוי באגים או מנסים להבין איך הקומפיילר עובד, לעיתים קרובות כדאי לקבל את ההרחבה של חומרה מסוימת בנקודה מסוימת בפייפליין (HLO, ‏ HLO מותאם, TritonIR או LLVM), עבור קלט HLO או StableHLO נתון.

hlo-opt תומך בשלבי פלט מרובים: PTX,‏ HLO אחרי אופטימיזציות, LLVM IR לפני אופטימיזציות או TritonIR. הקבוצה המדויקת של השלבים הנתמכים תלויה בפלטפורמה (למשל, PTX הוא ספציפי ל-NVIDIA), ואפשר לראות אותה באמצעות הפקודה ‎--list-stages:

hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx

אחרי שבוחרים שלב, המשתמש יכול לכתוב את התוצאה של ההמרה עבור פלטפורמה מסוימת לזרם נתונים מסוים:

hlo-opt --platform=cpu --stage=hlo input.hlo

הפקודה הזו תדפיס את ה-dump ל-stdout (או לקובץ נתון אם צוין -o).

קומפילציה של GPU ללא מכשיר

קומפילציה ללא מכשיר לא צריכה גישה ל-GPU. התכונה 'קומפילציה ללא מכשיר' מאפשרת לציין את מפרט ה-GPU בשורת הפקודה (--xla_gpu_target_config_filename) בשלבים שבהם נדרשת גישה ל-GPU, וכך לא צריך מכשיר GPU.

דוגמה: פלט PTX ללא גישה למכשיר GPU:

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo

המפרטים של מעבדי GPU פופולריים נשלחים עם המהדר, והקובץ שסופק הוא מחרוזת סריאליזציה של device_description.proto:

gpu_device_info {
  cuda_compute_capability {
    major: 8
    minor: 0
  }
  threads_per_block_limit: 1024
  threads_per_warp: 32
  shared_memory_per_block: 127152
  shared_memory_per_core: 65536
  threads_per_core_limit: 2048
  core_count: 6192
  fpus_per_core: 64
  block_dim_limit_x: 2147483647
  block_dim_limit_y: 65535
  block_dim_limit_z: 65535
  memory_bandwidth: 2039000000000
  l2_cache_size: 4194304
  clock_rate_ghz: 1.1105
  device_memory_size: 79050250240
}
platform_name: "CUDA"

מפרטים נוספים של GPU זמינים בכתובת /xla/tools/hlo_opt/gpu_specs

כוונון אוטומטי

לפעמים, תהליך ההרכבה כולל כוונון אוטומטי על סמך אוסף --stage. כדי שהקומפילציה ללא מכשיר תפעל, המשתמש צריך
להשבית את הכוונון האוטומטי באמצעות --xla_gpu_autotune_level=0
או
לטעון תוצאות קיימות של כוונון אוטומטי באמצעות --xla_gpu_load_autotune_results_from=<filename> (שהתקבלו באמצעות --xla_gpu_dump_autotune_results_to=<filename>).

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=gpu_specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo

קובץ הכוונון האוטומטי הוא סריאליזציה של טקסט של autotune_results.proto, עם דוגמה שנראית כך:

version: 3
results {
  device: "CUDA: 8.0, Cores: 108, GPU clock: 1.41 GHz, Memory bandwidth: 1555 GB/s, L2 cache: 40 MB"
  hlo: "{\n  tmp_0 = f16[1,16,17,3]{3,2,1,0} parameter(0)\n  tmp_1 = f16[16,51]{1,0} bitcast(f16[1,16,17,3]{3,2,1,0} tmp_0)\n  tmp_2 = s8[16,17,3]{2,1,0} parameter(1)\n  tmp_3 = s8[51,16]{0,1} bitcast(s8[16,17,3]{2,1,0} tmp_2)\n  tmp_4 = f16[51,16]{0,1} convert(s8[51,16]{0,1} tmp_3)\n  tmp_5 = f16[16,16]{1,0} dot(f16[16,51]{1,0} tmp_1, f16[51,16]{0,1} tmp_4), lhs_contracting_dims={1}, rhs_contracting_dims={0}\n  ROOT tmp_6 = f16[1,16,16]{2,1,0} bitcast(f16[16,16]{1,0} tmp_5)\n}"
  result {
    run_time {
      nanos: 31744
    }
    triton {
      block_m: 32
      block_n: 32
      block_k: 32
      split_k: 1
      num_stages: 1
      num_warps: 4
    }
  }
}

אפשר לבצע סריאליזציה של מסד הנתונים של הכוונון האוטומטי באמצעות XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>

‫[`hlo-opt`] פיתוח וניפוי באגים של כרטיסי HLO

# If you are working with hardware independent passes from the
# `xla/hlo/transforms/` directory, prefer light-weight version
# of the `hlo-opt` tool with fewer dependencies:

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

# Otherwise, for hardware independent and CPU, GPU passes use
# the same binary from "Compile HLO Modules" section above:

bazel run //xla/tools:hlo-opt -- [flags] <filename>

הכלי hlo-opt מאפשר להריץ מעברים בודדים בלי קשר לשלבי ההידור של הפלטפורמה הנתונה. הבידוד הזה עוזר להריץ במהירות מעברים במודול hlo של הקלט ולזהות את הגורם הראשוני לכשלים.

hlo-opt --passes=schedule-aware-collective-cse input.hlo

הכלי hlo-opt תומך גם ב-DebugOptions XLA_FLAGS.

hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo

משתמשים באפשרות--list-passes כדי לקבל את מחרוזת שם הכרטיס.

hlo-opt --list-passes

המשתמשים יכולים ליצור פייפליין משלהם בהתאמה אישית על ידי ציון יותר ממעבר אחד לאפשרות --passes.

hlo-opt --passes=pass1,pass2,pass3 input.hlo

עזרה בפיתוח של כרטיס New HLO

קודם כול, כותבים את הכרטיס.
רושמים את הכרטיס החדש במאגר הכרטיסים של הכלי hlo-opt.
```
RegisterPass<FooPass>(FooPassInputOptions)
```
בהתאם לסוג הכרטיס, בוחרים אחד מהמיקומים הבאים לרישום:
opt_lib.cc כרטיסים שלא תלויים בחומרה.
cpu_opt.cc כרטיסים ספציפיים למעבד.
gpu_opt.cc כרטיסים ספציפיים ל-GPU.
compiled_opt.cc העברות משותפות למעבד, ל-GPU ול-XPU.
אל תשכחו להוסיף תלות ב-build.

כדאי לכלול את רישום הכרטיס כחלק מה-PR(דוגמה) כדי שכל משתמשי hlo-opt יוכלו להשתמש בכרטיס.
בונים מחדש את הכלי hlo-opt, מאמתים את הרישום של הכרטיס באמצעות האפשרות --list-passes ואז מריצים את הכרטיס באמצעות האפשרות --passes.
```
$ hlo-opt --passes=foo-pass input.hlo
```
רוצים לכתוב בדיקות יחידה (unit testing) למעבר? פרטים נוספים זמינים בכתובת https://openxla.org/xla/test_hlo_passes.

העברת מדידה בזמן ריצה

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

time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo

‫[`hlo-opt`] המרת פורמטים של מודולים ב-HLO

# Use the light weight version of the `hlo-opt` tool.

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

המרת `HLO Text` -> `HLO Proto`

hlo-opt --emit-proto input.hlo

המרת `HLO Proto` או `HLO Proto Binary` -> `HLO Text`

hlo-opt input.pbtxt or input.pb

‫[`ptx-opt`] Compiler LLVM Module down to PTX

הכלי יפעיל צינור אופטימיזציה של LLVMIR ואז יקרא ל-CompileToPtx.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

הכלי יכול גם להציג את LLVMIR אחרי כל נתיב.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>

‫[`isolate_hlo`] בידוד הוראות HLO בעייתיות

אם יש לכם קובץ dump גדול של HLO ואתם חושדים שהוראה או קטע ספציפיים במודול HLO גורמים לקריסה, אתם יכולים להשתמש בכלי isolate_hlo.

הכלי הזה מחלץ הוראה אחת של HLO (ואת ההקשר הדרוש שלה) למודול HLO חדש וקטן יותר. האפשרות הזו שימושית מאוד ליצירת רפרודוסר מינימלי ברמת הקומפיילר.

חומרי עזר ומקור: הכלי isolate_hlo זמין במאגר OpenXLA. אפשר לראות את הספרייה xla/tools בקוד המקור של XLA.
שימוש: יוצרים את הכלי מעץ המקור של XLA. בדרך כלל צריך להזין קובץ מודול HLO (טקסט או פרוטו), את שם ההוראה לחילוץ ואת נתיב קובץ הפלט.
```
# Example usage after building XLA:
# ./build/tools/isolate_hlo --input=module.hlo --instruction_name=fusion.123 \
#   --output=isolated_fusion.123.hlo --input_format=txt --output_format=long_txt
```
אפשר לעיין בהודעת העזרה של הכלי (--help) כדי לראות את הסימונים הספציפיים ואפשרויות הפורמט.

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

‫[run_hlo_module] הפעלת מודולים של HLO

הרצת כמה מודולים של HLO

‫[multihost_hlo_runner] הרצת מודולים של HLO עם תמיכה ב-SPMD

הרצת מודולים מרובים של HLO עם תמיכה ב-SPMD

‫[hlo-opt] הידור מודול HLO