תהליך העבודה של פיתוח XLA בדרך כלל מתרכז סביב HLO IR, שמייצג חישוב פונקציונלי מבודד שמסופק למהדר. XLA מגיע עם מספר כלי שורת פקודה (מתוארים בהמשך) שצורכים HLO ומריצים אותו, או מספקים שלב הידור ביניים. השימוש בכלים כאלה הוא חיוני למחזור איטרציה מהיר של compile->modify->run
, כי ה-HLO גם ניתן להמחשה ויזואלית וגם לפריצה. לרוב, הדרך המהירה ביותר להבין ולתקן ביצועים או התנהגות של XLA היא הדרך המהירה ביותר להבין ולפתור אותה.
הדרך הקלה ביותר להשיג את ה-HLO לתוכנית שבה הידור באמצעות XLA היא בדרך כלל להשתמש במשתנה הסביבה XLA_FLAGS
:
XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
ששומרת את כל קובצי ה-HLO לפני האופטימיזציה בתיקייה שצוינה, יחד עם ארטיפקטים שימושיים רבים אחרים.
קטעי הקוד של HLO מריצים: run_hlo_module
הכלי run_hlo_module
פועל ב-HLO לפני אופטימיזציה, וכברירת מחדל, על מנת לרכז את החבילות, להפעיל אותן ולבצע השוואה עם ההטמעה של התרגום באמצעות פירוש הדבר. לדוגמה, הקריאה הרגילה להרצת קובץ קלט computation.hlo
ב-NVIDIA GPU ולבדיקת תקינותו היא:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
כמו בכל הכלים, אפשר להשתמש ב---help
כדי לקבל את רשימת האפשרויות המלאה.
הרצת קטעי קוד של HLO עם תמיכה ב-SPMD: multihost_hlo_runner
כלי ההרצה של Multihost HLO הוא כלי דומה מאוד, עם האזהרה שנתמכת ב-SPMD, כולל תקשורת בין מארחים. הפעלה טיפוסית נראית כך:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
הרצת כרטיסים/שלבים של הידור HLO: hlo-opt
במהלך ניפוי באגים או הבנת אופן הפעולה של המהדר, לרוב כדאי לקבל את ההרחבה לחומרה מסוימת בנקודה מסוימת בצינור עיבוד הנתונים (HLO, HLO שעבר אופטימיזציה, TritonIR או LLVM), לקלט HLO נתון (יציב).
ב-hlo-opt
יש תמיכה במספר שלבי פלט: PTX, HLO אחרי אופטימיזציות, LLVM IR לפני אופטימיזציות או TritonIR. הקבוצה המדויקת של השלבים הנתמכים תלויה בפלטפורמה (לדוגמה, PTX ספציפית ל-NVIDIA) וניתן לראות אותה באמצעות הפקודה --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
לאחר בחירת השלב, המשתמש יכול לכתוב לזרם נתון את תוצאת ההמרה של פלטפורמה נתונה:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
שתדפיס את קובץ ה-Dump ל-stdout (או לקובץ נתון אם צוין -o
).
שימוש ללא מכשיר
גישה ל-GPU אינה נחוצה לרוב ההידור. על ידי ציון מפרט GPU בשורת הפקודה, אנחנו יכולים לקבל, למשל פלט PTX ללא גישה למאיץ:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_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"
אם נדרשת כוונון אוטומטי, ייתכן שתיתקלו בבעיות בהידור ללא מכשירים. למזלנו, אנחנו יכולים לספק אותם גם בשורת הפקודה:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=gpu_specs/a100_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
קובץ הכוונון האוטומטי הוא סריאליזציה של טקסט ל-autotune_results.proto
, עם דוגמה שנראית כך:
version: 2
results {
device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
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_t=<myfile.pbtxt>
הרצת כרטיס מהדר יחיד
יש תמיכה גם בדגלים מ-XLA_FLAGS
, כך שאפשר להשתמש בכלי כדי לבדוק אם מפעילים כרטיס יחיד:
hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo