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

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

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

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

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

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

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

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

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

एसपीएमडी सहायता के साथ HLO स्निपेट चलाना: multihost_hlo_runner

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

hlo_runner_main  --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo

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

कंपाइलर के काम करने से जुड़ी जानकारी को डीबग करते समय या उसे समझने के दौरान, दिए गए (स्टेबल) एचएलओ इनपुट के लिए, पाइपलाइन में किसी खास पॉइंट (जैसे, HLO, ऑप्टिमाइज़ किया गया HLO, 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

जो डंप को 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