आकार और लेआउट

XLA Op का स्ट्रक्चर

एचएलओ का एक उदाहरण देखें:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

इसमें ये कॉम्पोनेंट शामिल होते हैं:

• ऑपरेशन का नाम: add.936
  • यह ऑपरेशन का यूनीक नाम है.
• शेप: bf16[8,1,1280,16384]
  • यह Op का आउटपुट शेप है. यहां dtype bf16 है और शेप [8,1,1280,16384] है.
• लेआउट (टाइलिंग के साथ): 3,2,0,1:T(8,128)(2,1)
  • इससे पता चलता है कि ऐरे को मेमोरी में कैसे सेव किया जाता है. 3,2,0,1 से मेमोरी में मौजूद ऐक्सिस के क्रम (जैसे, कॉलम मेजर, लाइन मेजर वगैरह) का पता चलता है. वहीं, T(8,128)(2,1) से इस्तेमाल की गई टाइलिंग और पैडिंग का पता चलता है.
  • लेआउट का इस्तेमाल करना ज़रूरी नहीं है. अगर यह विकल्प नहीं चुना जाता है, तो टाइलिंग नहीं होती है. साथ ही, डाइमेंशन को सबसे बड़े से सबसे छोटे क्रम में रखा जाता है.
• ऑपरेशन: add
  • कार्रवाई की जा रही है. यहां यह Add है. इसका ज़िक्र Op नाम में भी किया गया है.
• आर्ग्युमेंट: exponential.183, broadcast.3115
  • इस ऑपरेशन में दो आर्ग्युमेंट होते हैं, जिन्हें उनके यूनीक नामों के साथ तय किया जाता है.

आइए, एक और उदाहरण देखते हैं. यह एक फ़्यूज़न ऑप है:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

इसमें पहले बताए गए कॉम्पोनेंट के अलावा, ये कॉम्पोनेंट भी शामिल होते हैं:

• एट्रिब्यूट: kind और calls
  • इनसे, की जा रही कार्रवाई के बारे में ज़्यादा जानकारी मिलती है. इस मामले में: फ़्यूज़न.
• मेमोरी की जगह (मेमोरी स्पेस आइडेंटिफ़ायर): S(1)
  • इससे मेमोरी स्पेस/उस जगह का पता चलता है जहां ऐरे को सेव किया जाता है. S(1) यहां यह बताता है कि यह ऐरे, वीएमईएम (टीपीयू पर) में मौजूद है.
• इनपुट आर्ग्युमेंट %fusion.32 के लिए शेप और लेआउट की जानकारी

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

आकार

XLA ShapeProto प्रोटो (xla_data.proto) से, N-डाइमेंशनल ऐरे (संक्षेप में ऐरे) के डाइमेंशन की संख्या, साइज़, और डेटा टाइप के बारे में पता चलता है.

शब्दावली, नोटेशन, और कन्वेंशन

• किसी ऐरे के डाइमेंशन की सही संख्या, उन डाइमेंशन की संख्या होती है जिनका साइज़ 1 से ज़्यादा होता है.

• N डाइमेंशन वाले अरे के लिए, डाइमेंशन को 0 से लेकर N-1 तक नंबर दिया जाता है. डाइमेंशन का साइज़, शून्य से कम नहीं होना चाहिए. खास तौर पर, साइज़ 0 मान्य है. डाइमेंशन नंबर, सुविधा के लिए मनमाने लेबल होते हैं. इन डाइमेंशन नंबर के क्रम से, शेप के लेआउट में किसी खास माइनर/मेजर ऑर्डरिंग का मतलब नहीं है. लेआउट, LayoutProto प्रोटो के हिसाब से तय होता है.

• परंपरा के मुताबिक, डाइमेंशन को डाइमेंशन नंबर के बढ़ते क्रम में दिखाया जाता है. उदाहरण के लिए, [A x B x C] साइज़ वाले तीन डाइमेंशन वाले ऐरे के लिए, डाइमेंशन 0 का साइज़ A, डाइमेंशन 1 का साइज़ B, और डाइमेंशन 2 का साइज़ C होता है.

  XLA में कुछ यूटिलिटी, Python की तरह नेगेटिव इंडेक्सिंग की सुविधा भी देती हैं: डाइमेंशन-1, आखिरी डाइमेंशन होता है. यह N डाइमेंशनल ऐरे के लिए N-1 के बराबर होता है. उदाहरण के लिए, ऊपर बताए गए तीन डाइमेंशन वाले ऐरे के लिए, डाइमेंशन -1 का साइज़ C, डाइमेंशन -2 का साइज़ B वगैरह होता है.

• दो, तीन, और चार डाइमेंशन वाले ऐरे में अक्सर डाइमेंशन से जुड़े खास अक्षर होते हैं. उदाहरण के लिए, 2D कलेक्शन के लिए:

  • डाइमेंशन 0: y
  • डाइमेंशन 1: x

  3D ऐरे के लिए:

  • डाइमेंशन 0: z
  • डाइमेंशन 1: y
  • डाइमेंशन 2: x

  4D ऐरे के लिए:

  • डाइमेंशन 0: p
  • डाइमेंशन 1: z
  • डाइमेंशन 2: y
  • डाइमेंशन 3: x

• XLA API में मौजूद फ़ंक्शन, डाइमेंशन को डाइमेंशन नंबर के बढ़ते क्रम में लेते हैं. यह उस क्रम से मेल खाता है जिसका इस्तेमाल डाइमेंशन को initializer_list के तौर पर पास करते समय किया जाता है; जैसे कि

  ShapeUtil::MakeShape(F32, {A, B, C, D})

  इससे एक ऐसा शेप बनेगा जिसके डाइमेंशन साइज़ ऐरे में [A, B, C, D] क्रम शामिल होगा.

लेआउट

LayoutProto प्रोटो में बताया गया है कि मेमोरी में किसी ऐरे को कैसे दिखाया जाता है. इसमें ये फ़ील्ड शामिल होते हैं:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

माइनर से मेजर डाइमेंशन तक के क्रम में डेटा दिखाना

सिर्फ़ minor_to_major फ़ील्ड को भरना ज़रूरी है. इस फ़ील्ड में, किसी शेप में डाइमेंशन के छोटे से बड़े क्रम के बारे में बताया गया है. minor_to_major में मौजूद वैल्यू, ऐरे के डाइमेंशन (N डाइमेंशन वाले ऐरे के लिए 0 से N-1) का क्रम होती हैं. इसमें पहली वैल्यू सबसे छोटा डाइमेंशन होती है और आखिरी वैल्यू सबसे बड़ा डाइमेंशन होती है. सबसे कम अहमियत वाला डाइमेंशन वह डाइमेंशन होता है जो लीनियर मेमोरी में मौजूद ऐरे के एलिमेंट में बदलाव होने पर सबसे तेज़ी से बदलता है.

उदाहरण के लिए, [2 x 3] साइज़ का यह 2D ऐरे देखें:

a b c
d e f

यहां डाइमेंशन 0 का साइज़ 2 है और डाइमेंशन 1 का साइज़ 3 है. अगर लेआउट में minor_to_major फ़ील्ड [0, 1] है, तो डाइमेंशन 0 सबसे छोटा डाइमेंशन है और डाइमेंशन 1 सबसे बड़ा डाइमेंशन है. यह लीनियर मेमोरी में इस लेआउट के मुताबिक होता है:

a d b e c f

0 से लेकर N-1 तक के इस माइनर-टू-मेजर डाइमेंशन ऑर्डर को कॉलम-मेजर (दो डाइमेंशन वाले डेटा के लिए) कहा जाता है. डाइमेंशन के मोनोटोनिक क्रम के हिसाब से, कोड में इस लेआउट को "डिम 0 माइनर है" के तौर पर भी दिखाया जा सकता है.

दूसरी ओर, अगर लेआउट में minor_to_major फ़ील्ड [1, 0] है, तो लीनियर मेमोरी में लेआउट इस तरह होगा:

a b c d e f

N डाइमेंशन वाले ऐरे के लिए, N-1 से लेकर 0 तक के माइनर-टू-मेजर डाइमेंशन का क्रम, लाइन-मेजर (दो डाइमेंशन वाले ऐरे के लिए) के जैसा होता है. डाइमेंशन के मोनोटोनिक क्रम के हिसाब से, कोड में इस लेआउट को "डिम 0 मेजर है" के तौर पर भी दिखाया जा सकता है.

डिफ़ॉल्ट रूप से माइनर से मेजर तक के नोट का क्रम

नए बनाए गए शेप के लिए डिफ़ॉल्ट लेआउट "डाइमेंशन का क्रम सबसे बड़े से सबसे छोटे तक है" होता है. इसका मतलब है कि [N-1, ..., 0].

पैडिंग (जगह)

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

ऐरे में इंडेक्सिंग

index_util.h में मौजूद IndexUtil क्लास, शेप और लेआउट के हिसाब से मल्टीडाइमेंशनल इंडेक्स और लीनियर इंडेक्स के बीच बदलने के लिए यूटिलिटी उपलब्ध कराती है. मल्टीडाइमेंशनल इंडेक्स में, हर डाइमेंशन के लिए एक int64 इंडेक्स शामिल होता है. लीनियर इंडेक्स, एक int64 वैल्यू होती है. यह वैल्यू, ऐरे को सेव करने वाले बफ़र में इंडेक्स करती है. आकृतियों और लेआउट को आसानी से बनाने और उनमें बदलाव करने के लिए, एक ही डायरेक्ट्री में shape_util.h और layout_util.h देखें.

मेमोरी स्पेस आइडेंटिफ़ायर

एचएलओ में, हर ऐरे को मेमोरी स्पेस आइडेंटिफ़ायर के साथ एनोटेट किया जा सकता है. इसे S(n) के तौर पर लिखा जाता है.

• S(0) (अक्सर इसे शामिल नहीं किया जाता) का मतलब है कि डिवाइस में हाई बैंडविड्थ मेमोरी (एचबीएम) है.
• S(1), डिवाइस पर मौजूद वर्चुअल मेमोरी (वीएमईएम) को दिखाता है.
• S(2), S(3) वगैरह, डिवाइस के हिसाब से अतिरिक्त मेमोरी स्पेस से जुड़े होते हैं.
• S(5) होस्ट मेमोरी को दिखाता है.