एचएलओ से थंक तक

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

प्री-ऑप्टिमाइज़ेशन एचएलओ

हम प्री-ऑप्टिमाइज़ेशन एचएलओ मॉड्यूल से शुरुआत करते हैं. प्री-ऑप्टिमाइज़ेशन HLO में ऐसे ऑपरेशन (ops) शामिल नहीं होते जिन्हें XLA के लिए इंटरनल माना जाता है. जैसे, fusion या bitcast. इस चरण में, ऑप्स के पास कोई लेआउट नहीं होता. अगर उनके पास कोई लेआउट होता भी है, तो उसे अनदेखा कर दिया जाएगा. प्री-ऑप्टिमाइज़ेशन एचएलओ आम तौर पर, TensorFlow और JAX जैसे हायर-लेवल फ़्रेमवर्क से जनरेट होता है. XLA फ़्लैग -xla_dump_to का इस्तेमाल करने पर, ऑप्टिमाइज़ेशन से पहले के HLO को “before_optimizations.txt” फ़ाइल नाम वाले सफ़िक्स वाली फ़ाइल में डंप कर दिया जाता है.

एचएलओ मॉड्यूल को ऑप्टिमाइज़ करना

XLA:GPU पाइपलाइन, प्री-ऑप्टिमाइज़ेशन वाले HLO को ऑप्टिमाइज़ किए गए HLO में बदलती है. इसके लिए, यह कई पास चलाती है. पास को सिमैंटिक तरीके से ग्रुप किया जा सकता है. साथ ही, इन्हें इस क्रम में चलाया जा सकता है:

शार्डिंग से जुड़े पास

इसमें Shardy Partitioner जैसे पास या SPMD शार्डिंग के लिए पास शामिल हैं.

ऑप्टिमाइज़ेशन पास

इसमें कानूनी तौर पर मान्य किए जाने और आसान बनाने, दोनों तरह के पास शामिल हो सकते हैं.

कलेक्टिव ऑप्टिमाइज़ेशन पास

यह ऑप्टिमाइज़ेशन पास की तरह ही होता है, लेकिन यह कलेक्टिव ऑप्स पर फ़ोकस करता है.

लेआउट असाइनमेंट पास

हर एचएलओ ऑप को एक लेआउट असाइन किया जाता है, जो निर्देश के आकार का हिस्सा होता है. लेआउट से यह कंट्रोल होता है कि मेमोरी में टेंसर को फ़िज़िकली कैसे लेआउट किया जाता है.

लेआउट के साथ शेप का उदाहरण:

f32[10,20,30]{2,0,1}

तत्व के टाइप के बाद, शेप के लॉजिकल डाइमेंशन होते हैं. इसके बाद, लेआउट के क्रमपरिवर्तन को छोटे से बड़े क्रम में रखा जाता है. इस उदाहरण में, सबसे कम अहम डाइमेंशन 30 है, दूसरा सबसे कम अहम डाइमेंशन 10 है, और सबसे अहम डाइमेंशन 20 है.

लेआउट असाइनमेंट का मकसद, लालची रणनीति का इस्तेमाल करके, ज़रूरी फ़िज़िकल ट्रांसपोज़िशन की संख्या को कम करना है. यह कुछ लेआउट की सीमाओं से शुरू होता है.जैसे, cuDNN/cuBLAS लाइब्रेरी को लगातार डाइमेंशन की ज़रूरत होती है. इसके बाद, यह लेआउट को HLO ग्राफ़ में “नीचे” और फिर “ऊपर” की ओर ले जाता है. लेआउट के फैलने की प्रोसेस के आखिर में, कुछ निर्देशों के लेआउट अलग-अलग हो सकते हैं. इनमें से एक लेआउट, ऑपरेंड से मिला होता है और दूसरा लेआउट, उपयोगकर्ता से मिला होता है. इस टकराव को ठीक करने के लिए, copy HLO निर्देश डाला जाता है. इससे लेआउट, ऑपरेंड लेआउट से बदलकर निर्देश लेआउट हो जाता है.

लेआउट नॉर्मलाइज़ेशन पास

फ़िजिकल शेप का पता लगाना थोड़ा मुश्किल होता है. इसलिए, लेआउट नॉर्मलाइज़ेशन की मदद से शेप को इस तरह से फिर से लिखा जाता है कि वह डिफ़ॉल्ट लेआउट {rank-1, rank-2, …, 0} का इस्तेमाल करे. ऊपर दिए गए उदाहरण में, सामान्य किया गया आकार f32[20,10,30]{2,1,0} होगा. लेआउट बदलने वाले कॉपी ऑपरेशन को transpose और bitcast के कॉम्बिनेशन के तौर पर फिर से लिखा जाता है. फ़िलहाल, हम सभी ऑपरेशंस को सामान्य नहीं कर सकते. इसलिए, अब भी कुछ ऐसे ऑपरेशंस हैं जिनके लेआउट डिफ़ॉल्ट नहीं हैं. इनमें सबसे खास gather और dot हैं. सामान्य किए गए ऑप्स और सामान्य नहीं किए गए ऑप्स के बीच की सीमाओं पर, bitcast ऑप्स होंगे.ये ट्रांसपोज़ को दिखाते हैं. इसका मतलब है कि असाइन किए गए लेआउट के साथ ट्रांसपोज़, जो इसे फ़िज़िकली नो-ऑप बनाता है.

लेआउट नॉर्मलाइज़ेशन, कुछ इंप्लिसिट ट्रांसपोज़ को भी साफ़ तौर पर दिखाता है. यह इसलिए ज़रूरी है, क्योंकि कोडजन, साफ़ तौर पर दिखाए गए ट्रांसपोज़ को एक खास एमिटर की मदद से हैंडल कर सकता है. उदाहरण के लिए, रीशेप करने पर ऑपरेंड और नतीजे के बीच अलग-अलग फ़िज़िकल लेआउट हो सकता है. जैसे, अलग-अलग रैंक की वजह से. लेआउट नॉर्मलाइज़ेशन पास के हिस्से के तौर पर चलने वाला ReshapeDecomposer पास, रीशेप को transpose, रीशेप bitcast, और transpose के क्रम में बदल देता है.

लेआउट असाइनमेंट ऑप्टिमाइज़ेशन पास

यहां सबसे ज़रूरी पास, ट्राइटन फ़्यूज़न (GEMM फ़्यूज़न + Softmax/Layernorm फ़्यूज़न) या लाइब्रेरी कॉल को फिर से लिखना है. इस चरण में, ऑटोट्यूनिंग भी होती है. इसमें XLA, अलग-अलग एमिटर में से किसी एक को चुनता है. साथ ही, कनवोल्यूशन या डॉट के लिए सबसे अच्छा एल्गोरिदम चुनता है. इसके अलावा, यह Triton एमिटर की मदद से हैंडल किए गए फ़्यूज़न के लिए सबसे अच्छी टाइलिंग ढूंढता है.

फ़्यूज़न पास

दो मुख्य पास, PriorityFusion और Multi-Output फ़्यूज़न हैं.

PriorityFusion में, हम लागत मॉडल के हिसाब से फ़्यूज़न बनाते हैं. फ़्यूज़ करते समय, हम कई उपयोगकर्ताओं के साथ डुप्लीकेट किए गए ऑप्स की अनुमति देंगे. ऐसा तब होगा, जब ऑप को सभी उपयोगकर्ताओं के साथ फ़्यूज़ किया जा सकता हो. अगर मुमकिन हो, तो हम मौजूदा Triton Softmax फ़्यूज़न को बढ़ाने की अनुमति भी देंगे.

Multi-Output फ़्यूज़न एक अलग पास है. यह उन ऑप्स/फ़्यूज़न को फ़्यूज़ करने की अनुमति देता है जो एक ऑपरेंड शेयर करते हैं. यह डुप्लीकेट किए बिना, ऑपरेंड/ऑपरेंड फ़्यूज़न को उपयोगकर्ताओं में भी फ़्यूज़ कर सकता है. इसके लिए, अतिरिक्त आउटपुट जोड़े जाते हैं, ताकि फ़्यूज़ किए जाने वाले ऑप के अन्य उपयोगकर्ताओं को इन आउटपुट पर रीडायरेक्ट किया जा सके. इस पास को यह ध्यान रखना होगा कि HLO ग्राफ़ में साइकल न बन जाएं.

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

फ़्यूज़न के बाद कई पास

कलेक्शन से जुड़े कई पास (जैसे, उन्हें एसिंक में बदलना या कलेक्शन के लिए किसी खास क्रम को लागू करना).

आखिर में, हम CopyInsertion चलाते हैं. इससे यह पक्का किया जाता है कि इन-प्लेस ऑपरेशन, ऐसे डेटा को न बदलें जिसकी ज़रूरत अब भी कहीं और है.

ऑप्टिमाइज़ेशन के आखिर में, ऑप्टिमाइज़ किए गए एचएलओ को डंप कर दिया जाता है. इसके लिए, -xla_dump_to फ़्लैग का इस्तेमाल किया जाता है. इसे ऐसी फ़ाइल में डंप किया जाता है जिसके फ़ाइल नेम का सफ़िक्स "after_optimizations.txt" होता है. अगर आपको HLO को इंटरमीडिएट पास के बाद डंप करना है, तो -xla_dump_hlo_pass_re=.* फ़्लैग का इस्तेमाल करें. इससे HloModule में बदलाव होता है. इसके अलावा, इसे कुछ पास तक सीमित रखने के लिए, किसी रेगुलर एक्सप्रेशन का इस्तेमाल किया जा सकता है.

शेड्यूलिंग

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

इसके बाद, LatencyHidingScheduler पास चलता है और कंप्यूट-कम्यूनिकेशन ओवरलैप को ज़्यादा से ज़्यादा करने की कोशिश करता है. हालांकि, इससे मेमोरी का इस्तेमाल फिर से बढ़ सकता है.

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

बफ़र असाइनमेंट

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

1. HloDataflowAnalysis निर्देशों को HloValues (असल में लॉजिकल बफ़र) असाइन करता है. इन-प्लेस ऑपरेशंस के लिए, किसी ऑपरेंड के HloValue का फिर से इस्तेमाल किया जा सकता है. कोई ऑप, एक से ज़्यादा HloValue तय कर सकता है. उदाहरण के लिए, टपल के नतीजे के आकार के साथ.

2. HloAliasAnalysis, एलियासिंग कार्रवाइयों के लिए बफ़र को एक साथ लाने की कोशिश करता है. साथ ही, HloValue से HloBuffer तक मैपिंग का हिसाब लगाता है.

3. BufferAssignment, HloBuffers को बफ़र स्लाइस में मैप करता है. ऐसा इस तरह से किया जाता है कि एक ही बफ़र स्लाइस का इस्तेमाल, ओवरलैप होने वाली लाइफ़टाइम वाली अलग-अलग HloBuffers के लिए न किया जाए. जिन कार्रवाइयों के लिए एलियास का इस्तेमाल किया जा सकता है उनके लिए, थोड़ा ओवरलैप होना ठीक है. ऐसा हो सकता है कि एक HloBuffer के खत्म होने का समय, दूसरे HloBuffer के शुरू होने के समय से मेल खाए. -xla_dump_to फ़्लैग का इस्तेमाल करने पर, बफ़र असाइनमेंट के बारे में कुछ जानकारी को "after_optimizations-buffer-assignment.txt" नाम वाले सफ़िक्स वाली फ़ाइल में डंप कर दिया जाता है.

थंक

एचएलओ ग्राफ़ को ऑप्टिमाइज़ और शेड्यूल करने के बाद, उसे किसी खास बैकएंड (सीपीयू या जीपीयू) के लिए थंक के लीनियर सीक्वेंस में बदल दिया जाता है.

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

थंक उत्सर्जन

शेड्यूल किए गए HLO कंप्यूटेशन को थंक सीक्वेंस में बदलने की प्रोसेस को "थंक उत्सर्जन" कहा जाता है. इसे हर बैकएंड में मौजूद, खास एमिटर क्लास मैनेज करती है.

जीपीयू बैकएंड के लिए, इसे IrEmitterUnnested हैंडल करता है. EmitHloComputation, कंप्यूटेशन में शेड्यूल किए गए HLO निर्देशों की सूची को दोहराता है और उन्हें खास Emit... तरीके से भेजता है. जैसे, EmitFusion, EmitConvolutionThunk, EmitWhile). इनमें से हर तरीके से, सही Thunk ऑब्जेक्ट बनाए जाते हैं और उन्हें थंक सीक्वेंस में जोड़ा जाता है.

सीपीयू बैकएंड के लिए, ThunkEmitter यह भूमिका निभाता है और इसे इसी तरह से व्यवस्थित किया जाता है. फ़ाइनल ThunkSequence को CpuExecutable में एम्बेड किया गया है.

ध्यान दें कि एचएलओ मॉड्यूल के एंट्री कंप्यूटेशन में मौजूद हर निर्देश, थंक की फ़ाइनल सीक्वेंस में मौजूद थंक की संख्या के हिसाब से, किसी भी थंक (kTuple, kConstant, ..), एक या कई थंक (उदाहरण के लिए, क्रम से लगाने का निर्देश) से मेल खा सकता है.

कमांड बफ़र: जीपीयू पर एक्ज़ीक्यूशन को ऑप्टिमाइज़ करना

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

कमांड बफ़र को थंक सीक्वेंस में CommandBufferThunk के तौर पर दिखाया जाता है.

एमिटर, इस थंक को सीधे तौर पर एचएलओ निर्देशों से नहीं बनाता है. इसके बजाय, यह काम CommandBufferConversionPass करता है. यह ThunkSequence पर चलता है.

पास, एक साथ आने वाले उन सब-सीक्वेंस की पहचान करता है जिनमें एक जैसे थंक होते हैं. जैसे, KernelThunk और GemmThunk की एक सीरीज़. इसके बाद, मिले हुए सब-सीक्वेंस को एक CommandBufferThunk से बदल देता है. नया थंक, ओरिजनल थंक के लॉजिक को CommandBufferCmd ऑब्जेक्ट की सूची के तौर पर शामिल करता है. जब किसी जीपीयू स्ट्रीम पर पहली बार CommandBufferThunk लागू होता है, तो यह हार्डवेयर कमांड बफ़र में कमांड के क्रम को "रिकॉर्ड" करता है. इसके बाद, हर बार इसे चलाने पर, यह GPU को सिर्फ़ एक निर्देश देता है. इससे GPU, रिकॉर्ड किए गए क्रम को "फिर से चलाता" है. इससे हर कर्नल को लॉन्च करने के दौरान, सीपीयू के ओवरहेड से बचा जा सकता है.

एक्ज़ीक्यूटेबल

XLA कंपाइलेशन पाइपलाइन का फ़ाइनल प्रॉडक्ट, एक एक्ज़ीक्यूटेबल होता है. यह किसी प्लैटफ़ॉर्म के लिए खास तौर पर बनाया गया होता है और इसमें सभी ज़रूरी चीज़ें शामिल होती हैं. इस ऑब्जेक्ट में, कंपाइल किए गए प्रोग्राम को टारगेट डिवाइस (जैसे, सीपीयू या जीपीयू) पर चलाने के लिए ज़रूरी सभी जानकारी शामिल होती है. यह कंपाइलर और रनटाइम के बीच का ब्रिज है. PJRT जैसे मॉडर्न रनटाइम, थोड़े ज़्यादा लेवल के ऐब्स्ट्रैक्शन का इस्तेमाल करते हैं (PjRtExecutable देखें). हालांकि, ये आखिर में बैकएंड के हिसाब से एक्ज़ीक्यूटेबल को रैप करते हैं.

Executable में कई अहम जानकारी होती है, जो कंपाइल करने के दौरान जनरेट होती है. हालांकि, बैकएंड के हिसाब से कॉन्टेंट अलग-अलग हो सकता है. आम तौर पर, इनमें ये शामिल होते हैं:

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

• एक्ज़ीक्यूशन प्लान (ThunkSequence): यह रनटाइम लॉजिक का मुख्य हिस्सा होता है. यह थंक ऑब्जेक्ट का लीनियर सीक्वेंस है. हर थंक, काम की एक यूनिट को दिखाता है. जैसे, कर्नल लॉन्च करना, लाइब्रेरी फ़ंक्शन (जैसे, cuBLAS) को कॉल करना या कंट्रोल फ़्लो को मैनेज करना. रनटाइम, इस क्रम में दोहराकर प्रोग्राम को एक्ज़ीक्यूट करता है.

• मेमोरी लेआउट (BufferAssignment): यह मेटाडेटा का अहम हिस्सा है. इसे BufferAssigner बनाता है. इसमें कंप्यूटेशन के लिए पूरे मेमोरी लेआउट के बारे में बताया जाता है. इससे हर बफ़र का साइज़ तय होता है. साथ ही, यह तय होता है कि पैरामीटर, आउटपुट, और अस्थायी वैल्यू के लिए मेमोरी कैसे असाइन की जाती है और उसका दोबारा इस्तेमाल कैसे किया जाता है. रनटाइम इसका इस्तेमाल, डिवाइस की मेमोरी को असाइन करने और हर थंक को सही पॉइंटर पास करने के लिए करता है.

• (ज़रूरी नहीं) HLO मॉड्यूल: डीबग करने और प्रोफ़ाइलिंग के लिए, एक्ज़ीक्यूटेबल अक्सर उस फ़ाइनल, ऑप्टिमाइज़ किए गए HloModule का रेफ़रंस सेव रखता है जिससे इसे कंपाइल किया गया था.

हर बैकएंड के लिए, कंपाइलर फ़ाइनल एक्ज़ीक्यूटेबल बनाता है. कंपाइलर को लागू करने का RunBackend तरीका, कंपाइलेशन प्रोसेस का आखिरी चरण होता है. इसमें कंपाइल किए गए सभी आर्टफ़ैक्ट को एक्ज़ीक्यूटेबल ऑब्जेक्ट में पैकेज किया जाता है. GpuCompiler और CpuCompiler क्रमशः जीपीयू और सीपीयू को टारगेट करते हैं.

जब कोई उपयोगकर्ता किसी एक्ज़ीक्यूटेबल पर Execute... को कॉल करता है, तो रनटाइम, मेमोरी को असाइन करने के लिए BufferAssignment का इस्तेमाल करता है. इसके बाद, कंपाइल किए गए कोड का इस्तेमाल करके डिवाइस पर कार्रवाइयां शुरू करने के लिए, ThunkSequence को शुरू करता है.