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

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

תהליך הפיתוח של XLA מתמקד בדרך כלל ב-IR של HLO, שמייצג חישוב פונקציונלי מבודד שמסופק למהדר. XLA כולל כמה כלים של שורת הפקודה (מתוארים בהמשך) שמשתמשים ב-HLO ומריצים אותו, או מספקים שלב הידור ביניים. שימוש בכלים כאלה הוא שימושי מאוד למחזור חזרות (iteration) מהיר של compile->modify->run, כי אפשר להציג את ה-HLO באופן חזותי ולפרוץ אליו, ולרוב שינוי ורצף הפעלות של ה-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 ב-GPU של NVIDIA ולבדוק את תקינות הקובץ היא:

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

כמו בכל הכלים, אפשר להשתמש ב---help כדי לקבל את רשימת האפשרויות המלאה.

הפעלת קטעי קוד HLO עם תמיכה ב-SPMD: multihost_hlo_runner

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

הפעלה חוזרת של כמה HLO

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

$ hlo_runner_main /dump/*before_optimizations*

הרצת שלבים/תהליכים של הידור 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

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

שימוש ללא מכשיר

אין צורך בגישה ל-GPU ברוב תהליך ה-compilation, וניתן לציין מפרט GPU בשורת הפקודה כדי לקבל, למשל, פלט PTX בלי גישה למאיץ:

$ hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo

המפרטים של המעבדים הגרפיים הפופולריים נשלחים עם המהדר, והקובץ שסופק הוא שרשור של המחרוזת 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>

הפעלת סבב אחד של הידור

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

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