שימוש בכלים של XLA

תהליך העבודה של פיתוח 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