XLA के लिए नया बैकएंड डेवलप करना

यह गाइड उन सिस्टम इंजीनियर के लिए है जो XLA से अपने हार्डवेयर को बेहतर तरीके से टारगेट करने वाले प्रोग्राम जनरेट करना चाहते हैं. इस गाइड में सिलसिलेवार तरीके से जानकारी नहीं दी गई है. यह LLVM, Bazel, और XLA के बारे में जानकारी देती है.

XLA एक ऐब्स्ट्रैक्ट इंटरफ़ेस उपलब्ध कराता है, जिसे नया आर्किटेक्चर या ऐक्सेलरेटर इस्तेमाल करके, XLA एमएल प्रोग्राम आउटपुट चलाने के लिए बैकएंड बनाया जा सकता है. नए हार्डवेयर के लिए, PyTorch या TensorFlow जैसे फ़्रंटएंड फ़्रेमवर्क से सभी मौजूदा ऑपरेटर को लागू करने की तुलना में, XLA को फिर से टारगेट करना, काफ़ी आसान और बेहतर होना चाहिए.

ज़्यादातर तरीकों को इनमें से किसी एक स्थिति में लागू किया जाएगा:

  1. मौजूदा सीपीयू आर्किटेक्चर, फ़िलहाल XLA के साथ आधिकारिक तौर पर काम नहीं करता है. साथ ही, यह LLVM बैकएंड के साथ या उसके बिना भी काम नहीं करता है.
  2. मौजूदा एलएलवीएम बैकएंड के साथ बिना सीपीयू की तरह का हार्डवेयर.
  3. बिना सीपीयू की तरह का हार्डवेयर, जिसमें मौजूदा एलएलवीएम बैकएंड नहीं है.

पहली स्थिति: सीपीयू का मौजूदा आर्किटेक्चर, फ़िलहाल XLA पर आधिकारिक तौर पर काम नहीं करता है

इस स्थिति में, मौजूदा XLA सीपीयू बैकएंड पर गौर करके शुरुआत करें. XLA, LLVM का इस्तेमाल करके अलग-अलग सीपीयू को टारगेट करना आसान बनाता है, क्योंकि सीपीयू के लिए XLA बैकएंड के बीच मुख्य अंतर, LLVM से जनरेट किया गया कोड है.

अगर हार्डवेयर वेंडर के पास अपने हार्डवेयर के लिए एलएलवीएम बैकएंड है, तो बैकएंड को XLA के साथ बनाए गए एलएलवीएम से लिंक करना आसान है. JIT मोड में, XLA सीपीयू बैकएंड, होस्ट के सीपीयू के लिए कोड बनाता है. समय से पहले कंपाइल करने के लिए, xla::AotCompilationOptions टारगेट आर्किटेक्चर को कॉन्फ़िगर करने के लिए, एलएलवीएम ट्रिपल उपलब्ध करा सकता है.

अगर कोई एलएलवीएम बैकएंड मौजूद नहीं है, लेकिन किसी दूसरे तरह का कोड जनरेटर मौजूद है, तो मौजूदा ज़्यादातर सीपीयू बैकएंड का फिर से इस्तेमाल किया जा सकता है.

दूसरी स्थिति: मौजूदा एलएलवीएम बैकएंड के साथ, ऐसा हार्डवेयर जो सीपीयू जैसा नहीं है

मौजूदा xla::CPUCompiler और xla::GPUCompiler क्लास में कोई नया xla::Compiler लागू किया जा सकता है, क्योंकि इनसे पहले से ही LLVM IR का उत्सर्जन होता है. हार्डवेयर के टाइप के आधार पर, हो सकता है कि LLVM IR जनरेशन के कई पहलुओं में बदलाव करना पड़े. हालांकि, मौजूदा बैकएंड के साथ कई कोड शेयर किए जा सकते हैं.

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

तीसरी स्थिति: ऐसा हार्डवेयर जो सीपीयू (CPU) जैसा न हो, लेकिन मौजूदा LLVM बैकएंड के बिना

अगर एलएलवीएम का इस्तेमाल नहीं किया जा सकता, तो यह सबसे अच्छा विकल्प है कि आप अपनी पसंद के हार्डवेयर के लिए, XLA का नया बैकएंड लागू करें. इस विकल्प को चुनने में ज़्यादा मेहनत करनी पड़ती है. लागू की जाने वाली क्लास के बारे में यहां बताया गया है:

  • StreamExecutor: कई डिवाइसों के लिए, StreamExecutor के सभी तरीकों की ज़रूरत नहीं होती. ज़्यादा जानकारी के लिए, StreamExecutor लागू करने के मौजूदा तरीके देखें.
  • xla::Compiler: इस क्लास में, एचएलओ के कंप्यूटेशन के कंपाइलेशन को xla::Executable में शामिल किया जाता है.
  • xla::Executable: इस क्लास का इस्तेमाल, प्लैटफ़ॉर्म पर कंपाइल किए गए कंप्यूटेशन को लॉन्च करने के लिए किया जाता है.
  • xla::TransferManager: यह क्लास बैकएंड को यह सुविधा देती है कि वे दिए गए डिवाइस मेमोरी हैंडल से XLA लिटरल डेटा बनाने के लिए, प्लैटफ़ॉर्म के हिसाब से सिस्टम उपलब्ध करा सकें. दूसरे शब्दों में कहें, तो इससे होस्ट से डिवाइस पर और वापस डेटा ट्रांसफ़र होने में मदद मिलती है.