XLA टूलिंग का इस्तेमाल करना

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

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

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

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

एचएलओ स्निपेट चल रहे हैं: run_hlo_module

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

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

सभी टूल की तरह, --help का इस्तेमाल करके विकल्पों की पूरी सूची देखी जा सकती है.

SPMD की मदद से HLO स्निपेट चलाना: multihost_hlo_runner

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

मल्टी-एचएलओ रीप्ले

run_hlo_module और hlo_runner_main, दोनों के लिए एक से ज़्यादा मॉड्यूल के साथ invocate का इस्तेमाल किया जा सकता है. इससे, डंप डायरेक्ट्री में मौजूद सभी मॉड्यूल को फिर से चलाना आसान हो जाता है:

$ hlo_runner_main /dump/*before_optimizations*

एचएलओ कंपाइलेशन के पास/स्टेज चल रहे हैं: hlo-opt

कंपाइलर की प्रोसेस को डीबग करने या समझने के लिए, अक्सर किसी दिए गए (स्टैबल) एचएलओ इनपुट के लिए, पाइपलाइन (चाहे वह एचएलओ, ऑप्टिमाइज़ किया गया एचएलओ, TritonIR या LLVM हो) के किसी खास पॉइंट पर, किसी खास हार्डवेयर के लिए एक्सपैंशन पाना मददगार होता है.

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

यह डंप को स्टैंडर्ड आउटपुट (या -o की वैल्यू तय होने पर, किसी फ़ाइल में) में प्रिंट करेगा.

डिवाइस के बिना इस्तेमाल करना

ज़्यादातर कंपाइलेशन के लिए, जीपीयू का ऐक्सेस ज़रूरी नहीं होता. साथ ही, कमांड लाइन पर जीपीयू की जानकारी देकर, हमें ऐक्सेलरेटर के ऐक्सेस के बिना भी 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