एक्सएलए टूलिंग

XLA डेवलपमेंट का वर्कफ़्लो आम तौर पर HLO IR पर आधारित होता है. यह कंपाइलर को दिए गए अलग-अलग फ़ंक्शनल कंप्यूटेशन को दिखाता है. XLA में कई कमांड लाइन टूल होते हैं. इनके बारे में यहां बताया गया है. ये टूल, HLO का इस्तेमाल करते हैं और इसे चलाते हैं या कंपाइल करने की इंटरमीडिएट स्टेज उपलब्ध कराते हैं. ऐसे टूल का इस्तेमाल, तेज़ी से compile->modify->runइटरेशन साइकल के लिए बहुत ज़रूरी है. ऐसा इसलिए, क्योंकि एचएलओ को विज़ुअलाइज़ किया जा सकता है और इसमें बदलाव किया जा सकता है. साथ ही, इसे बार-बार बदलकर और चलाकर, अक्सर XLA के परफ़ॉर्मेंस या व्यवहार को समझने और उसे ठीक करने का सबसे तेज़ तरीका होता है.

XLA के साथ कंपाइल किए जा रहे प्रोग्राम के लिए, एचएलओ पाने का सबसे आसान तरीका आम तौर पर XLA_FLAGS एनवायरमेंट वैरिएबल का इस्तेमाल करना होता है:

$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point

यह फ़ोल्डर, ऑप्टिमाइज़ेशन से पहले की सभी HLO फ़ाइलों को सेव करता है. साथ ही, इसमें कई अन्य काम की चीज़ें भी होती हैं.

[`run_hlo_module`] Run HLO Modules

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

यह टूल, प्री-ऑप्टिमाइज़ेशन एचएलओ पर काम करता है. साथ ही, डिफ़ॉल्ट रूप से बंडल कंपाइल करता है, उन्हें चलाता है, और रेफ़रंस इंटरप्रेटर के साथ तुलना करता है.run_hlo_module उदाहरण के लिए, NVIDIA GPU पर किसी इनपुट फ़ाइल computation.hlo को चलाने और उसकी जांच करने के लिए, आम तौर पर इस कमांड का इस्तेमाल किया जाता है:

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

एक से ज़्यादा एचएलओ मॉड्यूल चलाना

run_hlo_module के लिए, एक से ज़्यादा एचएलओ मॉड्यूल के साथ इनवोकेशन की सुविधा उपलब्ध है. किसी डायरेक्ट्री से सभी एचएलओ मॉड्यूल चलाने के लिए:

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

[`multihost_hlo_runner`] Run HLO Modules With SPMD Support

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

मल्टीहोस्ट एचएलओ रनर, काफ़ी हद तक इसी तरह का टूल है. हालांकि, इसमें यह अंतर है कि यह एसपीएमडी के साथ काम करता है. इसमें क्रॉस होस्ट कम्यूनिकेशन भी शामिल है. ज़्यादा जानकारी के लिए, एक से ज़्यादा होस्ट वाले एचएलओ रनर देखें.

एसपीएमडी की सुविधा के साथ एक से ज़्यादा एचएलओ मॉड्यूल चलाना

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 या StableHLO इनपुट के लिए, पाइपलाइन में किसी खास पॉइंट पर किसी खास हार्डवेयर के लिए एक्सपैंशन की जानकारी मिले. यह जानकारी HLO, ऑप्टिमाइज़ किया गया HLO, TritonIR या LLVM के लिए हो सकती है.

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

इससे डंप को stdout पर प्रिंट किया जाएगा. अगर -o विकल्प दिया गया है, तो डंप को दी गई फ़ाइल में प्रिंट किया जाएगा.

जीपीयू के लिए डिवाइस के बिना कंपाइल करना

डिवाइस के बिना कंपाइल करने के लिए, जीपीयू का ऐक्सेस ज़रूरी नहीं है. डिवाइस के बिना कंपाइल करने की सुविधा, उन चरणों के लिए कमांड लाइन (--xla_gpu_target_config_filename) पर जीपीयू स्पेसिफ़िकेशन तय करने का तरीका उपलब्ध कराती है जहां जीपीयू का ऐक्सेस ज़रूरी होता है. इससे जीपीयू डिवाइस की ज़रूरत खत्म हो जाती है.

उदाहरण: जीपीयू डिवाइस का ऐक्सेस न होने पर PTX आउटपुट:

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_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`] एचएलओ पास डेवलपमेंट और डीबगिंग

# 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-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 विकल्प में एक से ज़्यादा पास तय करके, अपनी पसंद के मुताबिक कस्टम पाइपलाइन बना सकते हैं.

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

नए HLO पास को डेवलप करने में मदद करना

सबसे पहले, अपना पास लिखें.
नए पास को hlo-opt टूल पास रजिस्ट्री में रजिस्टर करें.
```
RegisterPass<FooPass>(FooPassInputOptions)
```
पास के टाइप के हिसाब से, रजिस्ट्रेशन के लिए इनमें से कोई एक जगह चुनें:
opt_lib.cc ऐसे पास जो हार्डवेयर पर निर्भर नहीं होते.
cpu_opt.cc सीपीयू के हिसाब से पास.
gpu_opt.cc जीपीयू के हिसाब से पास.
compiled_opt.cc सीपीयू, जीपीयू, और एक्सपीयू के लिए सामान्य पास.
बिल्ड डिपेंडेंसी जोड़ना न भूलें.

पास के रजिस्ट्रेशन को अपने पीआर(उदाहरण) में शामिल करें, ताकि पास का इस्तेमाल सभी hlo-opt उपयोगकर्ता कर सकें.
hlo-opt टूल को फिर से बनाएं. --list-passes विकल्प का इस्तेमाल करके, पास के रजिस्ट्रेशन की पुष्टि करें. इसके बाद, पास चलाने के लिए --passes विकल्प का इस्तेमाल करें.
```
$ hlo-opt --passes=foo-pass input.hlo
```
क्या आपको पास के लिए यूनिट टेस्ट लिखने हैं? ज़्यादा जानकारी के लिए, 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`] Convert HLO Module Formats

# 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 को डंप भी कर सकता है.

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

[run_hlo_module] Run HLO Modules

एक से ज़्यादा एचएलओ मॉड्यूल चलाना

[multihost_hlo_runner] Run HLO Modules With SPMD Support

एसपीएमडी की सुविधा के साथ एक से ज़्यादा एचएलओ मॉड्यूल चलाना

[hlo-opt] HLO मॉड्यूल कंपाइल करें