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

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

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

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

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

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

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

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

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

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

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

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

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

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

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