תהליך הפיתוח של 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. התכונה Deviceless Compilation (קומפילציה ללא מכשיר) מאפשרת לציין את מפרט ה-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
עזרה בפיתוח כרטיס חדש של 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רוצים לכתוב בדיקות יחידה למעבר? פרטים נוספים זמינים בכתובת 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>