XLA आर्किटेक्चर

XLA (Accelerated लीनियर ऐलजेब्रा) एक मशीन लर्निंग (एमएल) कंपाइलर है जो लीनियर ऐलजेब्रा को ऑप्टिमाइज़ करता है. इससे प्रोग्राम चलाने की स्पीड और मेमोरी के इस्तेमाल में सुधार होता है. इस पेज पर XLA कंपाइलर के मकसद और आर्किटेक्चर के बारे में खास जानकारी दी गई है.

मकसद

फ़िलहाल, XLA कई ML फ़्रेमवर्क फ़्रंटएंड के साथ काम करता है (इनमें PyTorch, TensorFlow, और JAX शामिल हैं) और यह OpenXLA प्रोजेक्ट का हिस्सा है. यह ML के लिए ओपन-सोर्स कंपाइलर टेक्नोलॉजी का एक ईकोसिस्टम है जिसे प्रमुख ML हार्डवेयर और सॉफ़्टवेयर संगठनों के साथ मिलकर बनाया गया है. OpenXLA प्रोजेक्ट बनने से पहले, XLA को TensorFlow प्रोजेक्ट में बनाया गया था. हालांकि, इसके बुनियादी मकसद पहले जैसे ही हैं:

  • ऐप्लिकेशन को प्रोसेस करने की स्पीड बढ़ाएं. कुछ समय तक चलने वाले ऑपरेशन के चलने में लगने वाला समय कम करने और रनटाइम से ओवरहेड खत्म करने के लिए, सबग्राफ़ कंपाइल करें. साथ ही, मेमोरी में लगने वाले समय को कम करने के लिए, पाइपलाइन वाले ऑपरेशन को फ़्यूज़ करें, और जाने-पहचाने टेंसर आकार के बारे में बताएं, ताकि ज़्यादा एग्रेसिव प्रॉपेगेशन को बढ़ावा दिया जा सके.

  • मेमोरी के इस्तेमाल को बेहतर बनाएं. मेमोरी के इस्तेमाल का विश्लेषण करके उसे शेड्यूल करें. इससे, बीच के लेवल पर स्टोरेज के कई बफ़र खत्म हो जाएंगे.

  • कस्टम ऑपरेशन पर निर्भरता कम करें. कम-लेवल के उन ऑपरेशन की परफ़ॉर्मेंस में सुधार करें जो मूल रूप से हाथ से जोड़े गए कस्टम ऑप की परफ़ॉर्मेंस से मेल खाते हैं. ऐसा करने से, कई कस्टम ऑपरेशन की ज़रूरत नहीं पड़ेगी.

  • पोर्टेबिलिटी में सुधार करें. नॉवेल हार्डवेयर के लिए एक नया बैकएंड लिखना आसान बनाएं. इससे, एमएल मॉडल का एक बड़ा हिस्सा उस हार्डवेयर पर बिना किसी बदलाव के चल सकेगा. यह नए हार्डवेयर के लिए, एक ही तरह के मोनोलिथिक ऑपरेशन को खास बनाने के तरीके से उलट है, जिसमें उन ऑपरेशन का इस्तेमाल करने के लिए मॉडल को फिर से लिखना पड़ता है.

यह सुविधा कैसे काम करती है

XLA कंपाइलर, StableHLO में तय किए गए एमएल फ़्रेमवर्क से मॉडल ग्राफ़ लेता है और उन्हें अलग-अलग आर्किटेक्चर के लिए मशीन निर्देशों में इकट्ठा करता है. StableHLO, वर्शन वाले ऑपरेशन सेट (HLO = हाई लेवल ऑपरेशन) के बारे में बताता है. यह एमएल फ़्रेमवर्क और कंपाइलर के बीच पोर्टेबिलिटी लेयर उपलब्ध कराता है.

आम तौर पर, मॉडल ग्राफ़ को टारगेट के लिए ऑप्टिमाइज़ किए गए एक्ज़ीक्यूटेबल प्रोसेस में बदलने वाली कंपाइलेशन प्रोसेस में ये चरण शामिल होते हैं:

  1. XLA, StableHLO ग्राफ़ पर कई बिल्ट-इन ऑप्टिमाइज़ेशन और विश्लेषण पास करता है जो टारगेट-इंडिपेंडेंट होते हैं, जैसे कि सीएसई, टारगेट-इंडिपेंडेंट ऑपरेशन फ़्यूज़न और कंप्यूटेशन के लिए रनटाइम मेमोरी तय करने के लिए बफ़र ऐनलिसिस. इस ऑप्टिमाइज़ेशन चरण के दौरान, XLA StableHLO भाषा को आंतरिक HLO भाषा में भी बदल देता है.

  2. इस बार, टारगेट से जुड़ी जानकारी और ज़रूरतों को ध्यान में रखते हुए, XLA ज़्यादा से ज़्यादा एचएलओ-लेवल ऑप्टिमाइज़ेशन के लिए बैकएंड को एचएलओ कंप्यूटेशन भेजता है. उदाहरण के लिए, जीपीयू बैकएंड, ऐसे फ़ंक्शन फ़्यूजन कर सकता है जो खास तौर पर जीपीयू प्रोग्रामिंग मॉडल के लिए फ़ायदेमंद हों. साथ ही, यह तय कर सकता है कि कंप्यूटेशन को स्ट्रीम में किस तरह बांटा जाए. इस चरण में, बैकएंड, कुछ खास कार्रवाइयों या उनके कॉम्बिनेशन को ऑप्टिमाइज़ की गई लाइब्रेरी कॉल से पैटर्न के मुताबिक भी मैच कर सकता है.

  3. इसके बाद, बैकएंड, टारगेट के हिसाब से कोड जनरेट करता है. XLA के साथ शामिल किए गए सीपीयू और जीपीयू बैकएंड, लो-लेवल आईआर, ऑप्टिमाइज़ेशन, और कोड जनरेशन के लिए LLVM का इस्तेमाल करते हैं. ये बैकएंड, एचएलओ कंप्यूटेशन को बेहतर तरीके से दिखाने के लिए, ज़रूरी एलएलवीएम आईआर उत्सर्जित करते हैं. इसके बाद, इस एलएलवीएम आईआर से नेटिव कोड निकालने के लिए एलएलवीएम शुरू करते हैं.

इस प्रोसेस में, XLA कंपाइलर इस तरह से मॉड्यूलर है कि कुछ नए एचडब्ल्यू आर्किटेक्चर को टारगेट करने के लिए किसी वैकल्पिक बैकएंड में स्लॉट देना आसान है. जीपीयू बैकएंड, LLVM NVPTX बैकएंड के ज़रिए, NVIDIA जीपीयू के साथ काम करता है. सीपीयू (CPU) बैकएंड, कई सीपीयू आईएसए के साथ काम करता है.