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

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

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

एक चरणबद्ध कदम के रूप में, हम वर्तमान में प्रक्रिया को बाधित करने की प्रक्रिया में हैं, lhlo बफर-असाइनमेंट चरण को पूरा करता है और lhlo बोली में एक MLIR मॉड्यूल के बजाय उत्सर्जन 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 जितना संभव हो उतना सक्षम के साथ अंत-से-अंत तक चलने योग्य मॉडल होने को प्राथमिकता देती है। इसका तात्पर्य यह है कि प्राथमिकता के आधार पर निम्नलिखित उद्देश्यों की सूची:

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

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

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

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

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

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

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

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

  • TFRT
    • (प्रो) उपयोग के लिए महान CUDA और HIP रैपर।
    • (प्रो) आसान पुस्तकालय कॉल (cuDNN, cuBLAS, cuFFT, आदि) को लागू करने के लिए, क्योंकि TFRT ops 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 के विपरीत) है।
  • कम से कम तत्व-वार ऑप्स के लिए गतिशील आकार का समर्थन आसान है।

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

  1. डिवाइस कोड LLVM IR के रूप में रहता है।
  2. पुराने एमिटर को LHLO की तरह बनाएं -> MLIR LLVM बोली:
    • (लागत) यदि हम अंततः मानक की ओर पलायन करना चाहते हैं तो फेंक-फेंकने का काम किया जाएगा।
    • (लाभ) यह आसान और यांत्रिक है। कम अवधि में किया जा सकता है।
    • (लाभ) यह (1) की तुलना में अधिक लाभ नहीं करता है।
  3. रिफ्लैक्टर पुराने उत्सर्जक LHLO की तरह होगा -> MLIR GPU + मानक + लूप:
    • (लागत) मानक के लिए मौजूदा उत्सर्जकों को उठाना कुछ चुनौतियों का सामना करता है। पॉइंटर्स और जीईपी को मेमरी और सब व्यू में बदलना होगा। सुनिश्चित करें कि 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) के लिए जाते हैं, लेकिन वृद्धिशील रूप से नहीं। ऑप द्वारा इसे करने का कोई तरीका नहीं है, क्योंकि सभी तत्व-उत्सर्जित ऑप्स एक ही ग्राफ में जुड़े हुए हैं। यह कार्य कई चालू बलों (कर्नेल जनरेटर, लीनलग) के एकीकरण बिंदु के रूप में भी काम कर सकता है।
  • अन्य सभी ऑप्स (1) के लिए जाते हैं। एक खिंचाव लक्ष्य के रूप में, उन्हें (3) या (4) में स्थानांतरित किया जा सकता है।

प्राथमिकता

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

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

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

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

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

विस्तृत रचना

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

यह कदम सभी मौजूदा 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 पर पहले से ही काफी काम किया जा रहा है, इसलिए काम की वास्तविक मात्रा को देखने की जरूरत है।