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

תהליך הפיתוח של XLA מתמקד בדרך כלל HLO IR, שמייצג פונקציות מבודדות תשומת לב שניתנה למהדר (compiler). 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

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

שידור חוזר של Multi-HLO

הפעלה עם כמה מודולים נתמכת גם ב-run_hlo_module וגם ב- hlo_runner_main, שלרוב נוח להפעיל מחדש את כל המודולים באשפה ספרייה:

$ hlo_runner_main /dump/*before_optimizations*

שלבים/שלבי הרצה של הידור HLO: hlo-opt

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

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

$ 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_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"

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

$ 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_t=<myfile.pbtxt>

הפעלת Pass Compiler Pass

יש תמיכה גם בדגלים של XLA_FLAGS, כך שאפשר להשתמש בכלי לבדיקה הפעלת כרטיס יחיד:

$ hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo