XLA Tooling

گردش کار توسعه 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 قبل از بهینه‌سازی را در پوشه مشخص شده، به همراه بسیاری از مصنوعات مفید دیگر ذخیره می‌کند.

[ run_hlo_module ] ماژول های HLO را اجرا کنید

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

ابزار run_hlo_module روی HLO پیش از بهینه‌سازی عمل می‌کند و به طور پیش‌فرض کامپایل، اجرا و مقایسه با پیاده‌سازی مفسر مرجع را بسته‌بندی می‌کند. به عنوان مثال، فراخوانی معمول برای اجرای یک فایل ورودی computation.hlo روی یک پردازنده گرافیکی 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، از جمله ارتباط بین میزبان‌ها، پشتیبانی می‌کند. برای جزئیات بیشتر به Multi-Host HLO Runner مراجعه کنید.

اجرای چندین ماژول 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

که خروجی را در خروجی استاندارد (یا در یک فایل داده شده اگر -o مشخص شده باشد) چاپ می‌کند.

کامپایل بدون نیاز به دستگاه برای پردازنده گرافیکی (GPU)

کامپایل بدون دستگاه نیازی به دسترسی به پردازنده گرافیکی (GPU) ندارد. کامپایل بدون دستگاه راهی برای مشخص کردن مشخصات پردازنده گرافیکی (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

مشخصات مربوط به پردازنده‌های گرافیکی محبوب به همراه کامپایلر ارائه می‌شوند و فایل ارائه شده، سریال‌سازی رشته‌ای از 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"

مشخصات بیشتر پردازنده گرافیکی در /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 سریال‌سازی متنی 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 Pass

# 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 to --passes خط لوله سفارشی خود را ایجاد کنند.

hlo-opt --passes=pass1,pass2,pass3 input.hlo

به توسعه‌ی HLO Pass جدید کمک کنید

  1. اول، گذرنامه خود را بنویسید.
  2. رمز عبور جدید را در رجیستری رمز عبور ابزار hlo-opt ثبت کنید.

    RegisterPass<FooPass>(FooPassInputOptions)
    

    بر اساس نوع کارت، یکی از مکان‌های زیر را برای ثبت نام انتخاب کنید:
    opt_lib.cc گذرگاه‌های مستقل از سخت‌افزار.
    cpu_opt.cc مسیرهای خاص CPU.
    gpu_opt.cc گذرگاه‌های مخصوص پردازنده گرافیکی.
    compiled_opt.cc مسیرهای مشترک بین CPU، GPU و XPU.
    فراموش نکنید که وابستگی ساخت (build dependency) را اضافه کنید.

    ثبت نام برای دریافت کارت عبور را به عنوان بخشی از PR ( مثال ) خود لحاظ کنید تا کارت عبور برای استفاده همه کاربران hlo-opt در دسترس باشد.

  3. ابزار hlo-opt را بازسازی کنید، ثبت نام موفقیت‌آمیز مجوز را با استفاده از گزینه --list-passes تأیید کنید و سپس از گزینه --passes برای اجرای مجوز استفاده کنید.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. برای مرحله‌ی پاس، تست واحد می‌نویسید؟ برای اطلاعات بیشتر به 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 ] ماژول کامپایلر LLVM تا PTX

این ابزار، خط لوله بهینه‌سازی LLVMIR را اجرا کرده و سپس CompileToPtx را فراخوانی می‌کند.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

این ابزار همچنین می‌تواند LLVMIR را بعد از هر مسیر، dump کند.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>