XLA (Accelerated लीनियर ऐलजेब्रा) एक मशीन लर्निंग (एमएल) कंपाइलर है, जो लीनियर ऐलजेब्रा को ऑप्टिमाइज़ करता है. इससे प्रोग्राम चलाने की स्पीड और मेमोरी के इस्तेमाल में सुधार होता है. इस पेज पर एक्सएलए कंपाइलर के मकसद और आर्किटेक्चर की खास जानकारी है.
मकसद
फ़िलहाल, XLA में ML फ़्रेमवर्क फ़्रंटएंड (PyTorch, TensorFlow, और JAX शामिल हैं) के साथ काम करता है और यह OpenXLA प्रोजेक्ट का हिस्सा है – एमएल के लिए ओपन-सोर्स कंपाइलर टेक्नोलॉजी का नेटवर्क जिसे बेहतरीन एमएल हार्डवेयर और सॉफ़्टवेयर संगठनों के साथ मिलकर डेवलप किया गया है. OpenXLA प्रोजेक्ट बनाने से पहले, XLA को TensorFlow प्रोजेक्ट में बनाया गया था. हालांकि, इसके बुनियादी मकसद पहले जैसे ही हैं:
एक्ज़ीक्यूशन की स्पीड बढ़ाएं. कुछ समय तक चलने वाले ऑपरेशन के काम करने में लगने वाले समय को कम करने और रनटाइम से ओवरहेड खत्म करने के लिए सबग्राफ़ कंपाइल करें. साथ ही, मेमोरी ओवरहेड को कम करने के लिए, पाइपलाइन वाले ऑपरेशन में फ़्यूज़ करें, और ज़्यादा तेज़ रफ़्तार से प्रोपेगेशन को बढ़ाने के लिए, जाने-पहचाने टेंसर आकार बनाएं.
मेमोरी के इस्तेमाल को बेहतर बनाएं. इंटरमीडिएट स्टोरेज बफ़र को खत्म करने के लिए, मेमोरी के इस्तेमाल का विश्लेषण करें और उसे शेड्यूल करें.
कस्टम ऑपरेशन पर भरोसा कम करें. कई कस्टम ऑपरेशन की ज़रूरत को खत्म करने के लिए, अपने-आप जुड़े हुए लो-लेवल के ऑपरेशन की परफ़ॉर्मेंस में सुधार करें, ताकि उन्हें मूल रूप से हाथ से जोड़े गए कस्टम ऑपरेशन की परफ़ॉर्मेंस से मैच किया जा सके.
पोर्टेबिलिटी बेहतर बनाएं. नॉवेल हार्डवेयर के लिए एक नया बैकएंड लिखना आसान बनाएं, ताकि एमएल मॉडल का एक बड़ा हिस्सा उस हार्डवेयर पर बिना किसी बदलाव के चल सके. वहीं दूसरी ओर, नए हार्डवेयर के लिए मोनोलिथिक ऑपरेशन में महारत हासिल करने के तरीके से बिलकुल उलटा होता है. ऐसे ऑपरेशन का इस्तेमाल करने के लिए, मॉडल को फिर से लिखने की ज़रूरत होती है.
यह कैसे काम करता है
XLA कंपाइलर, StableHLO में बताए गए एमएल फ़्रेमवर्क से मॉडल ग्राफ़ लेता है और उन्हें अलग-अलग आर्किटेक्चर के लिए मशीन से जुड़े निर्देशों में इकट्ठा करता है. StableHLO, वर्शन वाले ऑपरेशन सेट (एचएलओ = हाई लेवल ऑपरेशन) के बारे में बताता है. यह एमएल फ़्रेमवर्क और कंपाइलर के बीच पोर्टेबिलिटी लेयर उपलब्ध कराता है.
आम तौर पर, मॉडल ग्राफ़ को टारगेट के लिए ऑप्टिमाइज़ किए गए एक्ज़ीक्यूटेबल में बदलने वाली कंपाइलेशन प्रोसेस में ये चरण शामिल होते हैं:
XLA, StableHLO ग्राफ़ पर कई बिल्ट-इन ऑप्टिमाइज़ेशन और विश्लेषण पास परफ़ॉर्म करता है जो टारगेट-इंडिपेंडेंट होते हैं. जैसे, सीएसई, टारगेट-इंडिपेंडेंट ऑपरेशन फ़्यूज़न और कंप्यूटेशन के लिए रनटाइम मेमोरी तय करने के लिए बफ़र विश्लेषण. ऑप्टिमाइज़ेशन के इस चरण के दौरान, XLA भी StableHLO भाषा को एक अंदरूनी एचएलओ भाषा में बदल देता है.
XLA, HLO-लेवल के ऑप्टिमाइज़ेशन के लिए, बैकएंड को एचएलओ कंप्यूटेशन भेजता है. इस बार, ऐसा टारगेट से जुड़ी जानकारी और ज़रूरतों को ध्यान में रखते हुए किया जाता है. उदाहरण के लिए, जीपीयू बैकएंड, ऐसे ऑपरेशन फ़्यूज़न कर सकता है जो खास तौर पर जीपीयू प्रोग्रामिंग मॉडल के लिए फ़ायदेमंद होते हैं. साथ ही, यह तय कर सकते हैं कि कंप्यूटेशन को स्ट्रीम में कैसे बांटा जाए. इस चरण में, बैकएंड, ऑप्टिमाइज़ की गई लाइब्रेरी कॉल के लिए कुछ खास कार्रवाइयों या कॉम्बिनेशन के पैटर्न के हिसाब से भी काम कर सकता है.
इसके बाद, बैकएंड, टारगेट के हिसाब से कोड जनरेट करता है. XLA के साथ शामिल सीपीयू और जीपीयू के बैकएंड, लो-लेवल आईआर, ऑप्टिमाइज़ेशन, और कोड जनरेट करने के लिए LLVM का इस्तेमाल करते हैं. ये बैकएंड, एचएलओ कंप्यूटेशन को सही तरीके से दिखाने के लिए ज़रूरी एलएलवीएम आईआर उत्सर्जित करते हैं. इसके बाद, इस एलएलवीएम आईआर से नेटिव कोड निकालने के लिए, एलएलवीएम शुरू कर देते हैं.
इस प्रोसेस में, XLA कंपाइलर इस लिहाज़ से मॉड्यूलर है कि एचडब्ल्यू के कुछ नए आर्किटेक्चर को टारगेट करने के लिए, किसी अन्य बैकएंड में इसे आसानी से स्लॉट किया जा सकता है. फ़िलहाल, जीपीयू बैकएंड, LLVM NVPTX बैकएंड की मदद से NVIDIA जीपीयू के साथ काम करता है. सीपीयू (CPU) बैकएंड, कई सीपीयू (CPU) आईएसए के साथ काम करता है.