इस पेज का अनुवाद Cloud Translation API से किया गया है.
Switch to English

XLA के लिए MLIR कोडगैन

HloInstruction पर HloInstruction करता है और इस प्रतिनिधित्व पर कई अनुकूलन करता है, लक्षित उपकरणों के बीच इनमें से बहुत कुछ साझा करता है। कुछ बिंदु के रूप में एक रेखीय अनुसूची की गणना की जाती है और स्मृति बफर को प्रत्येक मूल्य को सांख्यिकीय रूप से सौंपा जाता है। इस अनुक्रम को ट्रेस करके और विशिष्ट रूप से डिवाइस के लिए उपयुक्त प्रतिनिधित्व उत्पन्न करने के लिए "एमिटर" को कॉल करके डिवाइस विशिष्ट कोडेन संचालित होता है (उदाहरण के लिए सीपीयू पर XLA प्रतिनियुक्ति पर एक एकल LLVM फ़ंक्शन, या GPU ऑपरेशन "थ्रेड्स" का एक सिक्वेंस) और संभवतः उत्पन्न PTX जब लक्ष्यीकरण जीपीयू)।

स्टेजिंग स्टेप के रूप में, हम वर्तमान में एक्सएलए के बफर-असाइनमेंट चरण को पूरा करने के बाद प्रक्रिया को इंटरसेप्ट करने की प्रक्रिया में हैं और lhlo बोली में एक एमएलआईआर मॉड्यूल के बजाय उत्सर्जन करते हैं। वहां से हम डिवाइस के आधार पर MLIR घटकों (मुख्य रूप से Linalg, affine, और GPU बोली) का उपयोग करके कोडगेन करते हैं।

नीचे कोडेन इनपुट के रूप में lhlo का उपयोग करके lhlo / GPU को lhlo रिकॉर्ड करने की योजना है।

कार्य

मेज़बान युक्ति
इनपुट प्रारूप HloInstruction * (कार्य 1) HloInstruction * (कार्य 1)
आउटपुट स्वरूप xla :: थंक (कार्य 2) LLVM IR (टास्क 3)
  • टास्क 1 होस्ट और डिवाइस इनपुट फॉर्मेट को HloInstruction * से LHLO तक बदल देता है।
  • टास्क 2 मेजबान के आउटपुट स्वरूप को थ्रक्स से "होस्ट के लिए कुछ लैंडिंग पैड" में बदलता है (नीचे देखें)।
  • टास्क 3 एलएलवीएम आईआर से एमएलआर के किसी रूप में डिवाइस आउटपुट को माइग्रेट करता है। यह इस परियोजना के लिए वैकल्पिक है, और विवरण के लिए "माइग्रेटिंग डिवाइस एलएलवीएम आईआर" अनुभाग देखें।

यह परियोजना LHLO-emitters जितना संभव हो उतना सक्षम के साथ अंत-से-अंत तक चलने योग्य मॉडल होने को प्राथमिकता देती है। इसका तात्पर्य यह है कि प्राथमिकता द्वारा उद्देश्यों की निम्नलिखित क्रम सूची:

  • एलएचए / जीपीयू को एलएचएलओ एमिटर के साथ चलाने योग्य बनाएं, मौजूदा थ्रोन्स और एमिटर अनमॉडिफाइड के साथ।
  • HHInstruction के संदर्भों को हटाएं LHLO में, केस द्वारा मामला:
    • एक MLIR- आधारित एमिटर (जैसे Linalg), या के लिए एक विरासत उत्सर्जक स्विच करें
    • MLIR प्रतिनिधित्व लेने के लिए यंत्रवत् रूप से अनुवाद करें (GPU बोली के साथ मानक पर माइग्रेट करें)।

माइग्रेटिंग थ्रो (टास्क 2)

xla :: gpu :: Thunk एक डेटा संरचना है जो:

  • मेजबान से बुलाया जा सकता है (xla :: gpu :: Thunk :: ExecuteOnStream ())।
  • अपने उपवर्गों में विभिन्न डेटा वहन करती है।
  • BufferAllocation के साथ सहभागिता :: स्लाइस और StreamExecutor।
  • गुठली लॉन्च की
  • सभी रनटाइम लाइब्रेरी में कॉल करें।

उस की लागत में शामिल हैं:

  • ऑप-विशिष्ट कॉन्फ़िगरेशन डेटा (उदाहरण के लिए कॉन्फोल्यूशन कॉन्फ़िगरेशन) का प्रतिनिधित्व करना।
  • ओप आकार और ऑपरेंड आकृतियों को माइग्रेट करना।
  • थ्रक्स (जबकि, स्थिति, आदि) के एक पेड़ का प्रतिनिधित्व करना।

माइग्रेशन का काम LHLO / एमिटर माइग्रेशन से स्वतंत्र है। सीमित संसाधनों के तहत, यह एलएचएलओ / एमिटर माइग्रेशन के पीछे प्राथमिकता है।

हमारे पास LHLO से होस्ट-साइड वाले हिस्से को कम करने के कई विकल्प हैं:

  • TFRT
    • (प्रो) उपयोग के लिए महान CUDA और HIP रैपर।
    • (Pro) पुस्तकालय कॉल (cuDNN, cuBLAS, cuFFT, आदि) को लागू करना आसान है, क्योंकि TFRT ऑप्स की व्याख्या C ++ कोड द्वारा की जाती है।
    • (Con) मेजबान पक्ष विकास के अधीन है और परीक्षण नहीं किया गया है।
  • सीपीयू कोड
    • (प्रो) महान कम क्षमता। कुछ लूप और शर्तें बनाएं और यह पूरा हो गया है।
    • (Con) GPUDialect अभी तक चेन / स्ट्रीम / एसिंक्रोनसिटी / डिवाइस आवंटन को मॉडल नहीं करता है।
    • (Con) CUDA / HIP रनटाइम सपोर्ट न्यूनतम (टूलकिट पथ, संस्करण, डायनेमिक लोडिंग, आदि) है।
  • मौजूदा (व्याख्या) XLA रनटाइम

निर्णय: TFRT को अपनाएं, लेकिन TFRT में जटिंग सीपीयू कोड का भी समर्थन करें।

माइग्रेटिंग डिवाइस LLVM IR (टास्क 3)

एक तात्विक उत्सर्जक तत्व द्वारा तत्व को भरकर लक्ष्य ऑप उत्पन्न करता है। प्रत्येक आउटपुट तत्व ऑपरेंड से तत्वों के एक सेट पर निर्भर करता है। सभी तत्वों को गतिशील सूचकांकों के साथ बफर के संयोजन द्वारा वर्णित किया गया है। यह लगभग सभी "गणित" ऑप्स का वर्णन करने के लिए पर्याप्त है, लेकिन प्रदर्शन कारणों से केवल "गणित" ऑप्स का एक बड़ा उपसमुच्चय सीधे (Cpu | Gpu) में लागू किया जाता है ElementalIrEmitter।

ElementalIrEmitter उस में अद्वितीय है:

  • कोड का एक बड़ा हिस्सा XLA / GPU और CPU के बीच साझा किया जाता है।
  • यह सभी तत्व-वार ऑप्स सहित मॉडल में देखे गए ऑप्स के एक बड़े हिस्से का प्रतिनिधित्व करता है।
  • अधिकांश फ्यूजन पूरी तरह से ElementalIrEmitter पर निर्भर करते हैं।
  • यह संरचनात्मक रूप से सरल है, क्योंकि यह op तत्वों और ऑपरेंड तत्वों के बीच डेटा निर्भरता DAG का वर्णन करता है।
  • यह ज्यादातर पोर्टेबल और उच्च-स्तरीय (उदाहरण के लिए GPU kReduce और GPU kCopy के विपरीत) है।
  • कम से कम तत्व-वार ऑप्स के लिए गतिशील आकार का समर्थन आसान है।

अब सभी ऑप्स के लिए, मौलिक रूप से उत्सर्जित या नहीं, प्रत्येक XLA ऑप की अंतिम स्थिति के कई स्वाद हैं:

  1. डिवाइस कोड LLVM IR के रूप में रहता है।
  2. पुराने एमिटर को LHLO की तरह बनाएं -> MLIR LLVM बोली:
    • (लागत) यदि हम अंततः मानक की ओर पलायन करना चाहते हैं, तो फेंक-फेंकने का काम किया जाएगा।
    • (लाभ) यह आसान और यांत्रिक है। कम अवधि में किया जा सकता है।
    • (लाभ) यह (1) की तुलना में अधिक लाभ नहीं करता है।
  3. रिफ्लैक्टर पुराने उत्सर्जक जैसे LHLO -> MLIR GPU + मानक + लूप:
    • (लागत) मानक के लिए मौजूदा उत्सर्जकों को उठाना कुछ चुनौतियों का सामना करता है। पॉइंटर्स और GEPs को मेमरी और सब व्यू में परिवर्तित करने की आवश्यकता है। सुनिश्चित करें कि amdgpu पूर्णता एक और है।
    • (लागत) XLA / GPU एलएलवीएम मेटाडेटा पर बहुत निर्भर करता है:
      • ब्लॉक / थ्रेड इंडेक्स के लिए range
      • align , dereferenceable , invariant.load , alias.scope , लोड / स्टोर्स के लिए noalias
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable अनुक्रमिक छोरों के लिए उपयुक्त।
    • (लाभ) दीर्घकालिक हो सकता है। अधिक पोर्टेबल।
  4. रिफ्लैक्टर पुराने उत्सर्जक LHLO -> लिनाल्ग, और नए लीनिग उत्सर्जक लिखते हैं
    • (लागत) यह मामला है। पिछले विकल्पों की तुलना में, XLA के प्रदर्शन से मेल खाने वाले एक नए कार्यान्वयन को बेंचमार्क के माध्यम से जाने की आवश्यकता है <-> वर्कफ़्लो का अनुकूलन, जो कुछ ऑप्स के लिए एक महत्वपूर्ण लागत हो सकता है।
    • (लाभ) एकीकृत स्टैक; सामुदायिक समर्थन; पोर्टेबिलिटी; अधिक अनुकूलन क्षमता।

निष्कर्ष:

  • (2) के लिए मत जाओ। (1) या (3) सिर्फ (2) से बेहतर हैं। (2) की लागत (1) से अधिक है, क्योंकि इसके लिए बहुत सारे यांत्रिक रीफैक्टरिंग की आवश्यकता होती है। (1) के साथ हम अभी भी एमएलए एमिटर लेने के लिए एक्सएलए को सक्षम करने के लक्ष्य को प्राप्त कर सकते हैं। यह LHLO -> LLVM IR -> रन लेगेसी डिवाइस एमिटर कर रहा है।
  • ElementalIrEmitter ops (4) के लिए जाते हैं, लेकिन वृद्धिशील रूप से नहीं। ऑप द्वारा इसे करने का कोई तरीका नहीं है, क्योंकि सभी तत्व-उत्सर्जित ऑप्स एक ही ग्राफ में जुड़े हुए हैं। यह कार्य कई चालू बलों (xla / सेवा / mlir_gpu, कर्नेल जनरेटर, Linalg) के एकीकरण बिंदु के रूप में भी काम कर सकता है।
  • अन्य सभी ऑप्स (1) के लिए जाते हैं। एक खिंचाव लक्ष्य के रूप में, उन्हें (3) या (4) में स्थानांतरित किया जा सकता है।

प्राथमिकता

जबकि उपर्युक्त सभी तीन कार्य समानांतर हैं, सीमित संसाधनों के तहत उन्हें क्रमबद्ध किया जाना है। प्राथमिकताकरण प्रत्येक कार्य को पूरा करने के लिए दृश्यमान परिणामों पर केंद्रित है।

प्राथमिकता है: टास्क 1 (विरासत उत्सर्जक के लिए एलएचएलओ)> टास्क 2 (थ्रो)> टास्क 3 (एमएलआईआर एमिटर)।

टास्क 1 के अंत तक, XLA के उपयोगकर्ता एक LHLO (जैसे कर्नेल जनरेटर) उत्पन्न कर सकते हैं और उन्हें निष्पादित कर सकते हैं। संकलन प्रारूप क्रमबद्ध MLIR नहीं होगा।

टास्क 2 के अंत तक, एलएचएलओ उचित, सीरियल करने योग्य एमएलआईआर को कम करता है। यह ऑफ़लाइन संकलन सक्षम करता है।

टास्क 3 के अंत तक, सभी XLA एमिटर इसके कार्यान्वयन में MLIR- आधारित हैं।

विस्तृत रचना

चरण 1: (टास्क 1) पूरा एलएचएलओ बनाएं और लीगेसी एमिटर को एलएचएलओ बनाएं

यह कदम सभी मौजूदा XLA / GPU एमिटर को MLIR ऑप्स के साथ इंटरैक्ट करता है। यह कदम शुद्ध रिफैक्टरिंग और एनएफसी है।

यह कदम ज्यादातर यांत्रिक है, लेकिन यह एक अनावश्यक HloComputation और LHLO के बीच निम्नलिखित विसंगतियों को ध्यान देने योग्य है:

  • प्रत्येक HloInstruction की अपने ऑपरेंड्स (डेटा-प्रवाह DAG) तक सीधी पहुंच है। इसके विपरीत, प्रत्येक एलएचएलओ ओपी के पास केवल अपने ऑपरेंड बफ़र्स (ऑप्स और बफ़र्स के बीच एक द्विदलीय) तक पहुंच है। एलएचएलओ ऑप्स को अपने ऑपरेंड ऑप्स का उपयोग करने के लिए उपयोग-डेफ चेन से गुजरना पड़ता है।
  • अनुभवहीन विरासत अनुभवजन्य रूप से लगभग कभी भी अपने ऑपरेंड तक नहीं पहुंचती है। एकमात्र अपवाद kReduce है।
  • केवल उन स्लाइस प्राप्त करने के लिए अप्रयुक्त विरासत एमिटर एक्सेस बफ़रएसिग्नेशन उपयोग करती है, डेटाफ़्लो_नालिसिस () या अलियास_नालिसिस () जैसी ऑक्स डेटा संरचनाओं तक पहुँचने के लिए नहीं। llvm_ir स्लाइस जानकारी के आधार पर अपना स्वयं का अलियास_नालिसिस () बनाता है।

निष्कर्ष यह है कि एलएचएलओ को बड़ी परेशानी के बिना सही फिट होना चाहिए।

चरण 2: (वैकल्पिक) प्रोफाइलिंग सपोर्ट

यह कदम तभी आवश्यक है जब हम XLA Thunk लॉजिक में से कुछ को त्यागना शुरू करें (अगला चरण देखें)।

वास्तव में किसी भी MLIR- आधारित एमिटर को चालू करने से पहले, हमें MLIR- आधारित एमिटर के लिए प्रोफाइलिंग की आवश्यकता होती है।

वर्तमान में XLA, StreamExecutor के टाइमर में कॉल करके अपनी प्रोफाइलिंग करता है। हुड के नीचे टाइमर कर्नेल लॉन्च से पहले और बाद में दो घटनाओं को सम्मिलित करता है, और इन दो घटनाओं के बीच सिंक समय को मापता है।

MLIR में प्रोफाइलिंग का समर्थन करने के लिए लगभग तीन दृष्टिकोण हैं:

  • किसी प्रॉइलर को एंड-टू-एंड चलाएं
  • LHLO में प्रत्येक सेशन के लिए एक प्रोफाइल ऑप्‍शन जोड़ें, एक इंजेक्टेड प्रोफाइलर का उपयोग करें।

"एंड-टू-एंड" दृष्टिकोण MLIR के लिए पारदर्शी है, लेकिन उसी समस्या से ग्रस्त है जो XLA का उपयोग पहली जगह में नहीं करता है: एक प्रोफाइलर (nvprof / ...) द्वारा एकत्र किए गए पुस्तकालय कॉल आसानी से HLO से संबंधित नहीं हो सकते हैं ऑप्स। उदाहरण के लिए, cuDNN प्रत्येक HLO के लिए कई गुठली लॉन्च करता है, और यह बताना मुश्किल है कि कौन सी गुठली किस HLO के अनुरूप है।

"इंजेक्टेड प्रोफाइलर" दृष्टिकोण की आवश्यकता है:

  • एक पैरामीटर के रूप में एक प्रोफाइलर लेने के लिए LHLO।
  • प्रत्येक सेशन से पहले और बाद में profile.start / profile.end डालें।
  • उस प्रोफ़ाइल से एक पास प्रोफ़ाइल, {प्रारंभ, अंत} से एक C ++ कार्यान्वयन के लिए।

सटीक प्रोफाइलिंग आसानी से एमएलआईआर-जनित ऑप्स के लिए नहीं की जा सकती है, क्योंकि:

  • MLIR में टाइमर नहीं है, न ही यह TFRT / StreamExecutor पर निर्भर करता है।
  • MLIR जटिल मापदंडों के साथ आसानी से C फ़ंक्शन में कॉल नहीं करता है।

चरण 3: (कार्य 2) माइग्रेटिंग थ्रो

नोट के रूप में, मोटे तौर पर तीन प्रकार के होते हैं:

  • कर्नेलटहंक, जो कर्नेल लॉन्च करता है।
  • कंट्रोल फ्लो थ्रक्स, जिसमें होस्ट कंट्रोल फ्लो लॉजिक (सशर्त, जबकि, सीक्वेंस) और बॉडी कर्नेल लॉन्च होते हैं।
  • लाइब्रेरी थ्रक्स: cuDNN, cuBLAS, cuFFT, NCCL, आदि।

योजना है:

  • थ्रॉक्स (डी) को सीरियल बनाने योग्य बनाएं।
  • TFRT को ऐसी स्थिति में सुधार करने में मदद करें जहां यह इन शब्दार्थों का समर्थन कर सकता है।
  • जैसे ही राज्य में सुधार होता है, अलग-अलग थ्रेड को वृहद रूप से माइग्रेट करें।

ये एक्शन आइटम केवल आंशिक रूप से ऑर्डर किए गए हैं। वास्तविक निष्पादन आदेश / इंजीनियरिंग समानता का मूल्यांकन किया जाता है क्योंकि यह जाता है।

चरण 4: (टास्क 3) माइग्रेटेड एलिमेंटलइमित्र

एक बार प्रोफाइलिंग तैयार हो जाने के बाद, हम MLIRIrEmitter- आधारित एमिटर को MLIR में पूरा और ट्यून कर सकते हैं। फिर हम उन्हें डिफ़ॉल्ट रूप से चालू करते हैं, यह मानते हुए कि ये सभी एमएलआईआर-आधारित एमिटर एक धारा का उपयोग करते हैं।

ध्यान दें कि XLA / CPU के ElementalIrEmitter को माइग्रेट करना फायदेमंद है, क्योंकि वे कोड का एक बड़ा हिस्सा साझा करते हैं।

सभी बेंचमार्किंग और प्रदर्शन शिकार के साथ (TODO: प्रदर्शन समता को परिभाषित), हम नए MLIR- आधारित तात्विक उत्सर्जक को चालू करते हैं, और विरासत ElementalIrEmitter को हटाते हैं।

यह कदम बाद के प्रवास के लिए आसान संलयन संक्रमण (नेस्टेड ऑप्स) भी प्रदान करता है।

चरण 5: मल्टी-स्ट्रीम समर्थन या ड्रॉप

जब तक हम MLIR में इसका समर्थन नहीं करते, या हम इस सुविधा को नहीं छोड़ते, हम कुछ उत्सर्जकों को हटा नहीं सकते। यह MLIR में अपेक्षाकृत बड़ी मात्रा में काम है और XLA के लिए थोड़ी मात्रा में लाभ है। हमें बहु-धारा XLA / GPU उपयोगकर्ताओं के वर्तमान उपयोगकर्ताओं की जांच करनी चाहिए, और यदि उचित हो तो इस सुविधा को हटाने का प्रयास करें।

चरण 6: (कार्य 3) माइग्रेटेड डिवाइस ऑप्स

यह चरण सभी अनावश्यक ऑप्स को माइग्रेट करता है, फिर हम सभी अनावश्यक उत्सर्जकों को हटा सकते हैं।

यह kCopy और kReduce के लिए फिर से लिखना / रिफलेक्टर पर कॉल करता है। kReduce पर पहले से ही काफी काम किया जा रहा है, इसलिए काम की वास्तविक मात्रा को देखने की जरूरत है।