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