با استفاده از ابزار 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 و بررسی صحت آن عبارت است از:

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

مانند همه ابزارها، می توان --help برای به دست آوردن لیست کامل گزینه ها استفاده کرد.

اجرای قطعه HLO با پشتیبانی SPMD: multihost_hlo_runner

Multihost HLO runner یک ابزار بسیار مشابه است، با این نکته که از 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-stags قابل مشاهده است:

$ 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

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