استخدام أدوات 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

يُعد برنامج 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

يطبع ملف التفريغ إلى stdout (أو إلى ملف معيّن في حال تحديد -o).

الاستخدام بدون أجهزة

ليس من الضروري الوصول إلى وحدة معالجة الرسومات في معظم مرحلة تجميع البيانات، ويمكنك من خلال تحديد مواصفات وحدة معالجة الرسومات في سطر الأوامر الحصول على سبيل المثال، مخرجات 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_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