'sdy' बोली

Shardy (SDY) डायлект, ऐक्सिस पर आधारित टेंसर का बंटवारा करने के तरीके के बारे में बताता है. साथ ही, टेंसर में बंटवारे को अटैच करने के लिए अतिरिक्त एपीआई कॉम्पोनेंट भी बताता है.

ऑपरेशंस

sdy.constant (sdy::ConstantOp)

कंटेंट के लिए लगातार काम करना

यह किसी कॉन्स्टेंट value से output टेंसर बनाता है.

यहां देखें: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

उदाहरण:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

विशेषताएं: AlwaysSpeculatableImplTrait

इंटरफ़ेस: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

इफ़ेक्ट: MemoryEffects::Effect{}

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
value::mlir::ElementsAttrकॉन्सटेंट वेक्टर/टेंसर एट्रिब्यूट

नतीजे:

नतीजा ब्यौरा
output किसी भी तरह की वैल्यू का टेंसर

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

डेटा फ़्लो एज ऑपरेशन.

सिंटैक्स:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

किसी ऑपरेटर X का डेटा फ़्लो एज, सोर्स के एक सेट (हर सोर्स, X का ऑपरेंड या X के ब्लॉक टर्मिनेटर का ऑपरेंड होता है) और टारगेट के एक सेट (हर टारगेट, X का नतीजा या X का ब्लॉक आर्ग्युमेंट होता है) के बीच एक ब्रिज तय करता है. ऐसा इसलिए किया जाता है, ताकि सभी सोर्स और टारगेट को एक ही तरह से शर्ड किया जा सके.

किसी ऑप में डेटा फ़्लो के ऐसे कई किनारे हो सकते हैं जो एक-दूसरे के ऑर्थोगोनल होते हैं.

उदाहरण के लिए:

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

इस while op में n डेटा फ़्लो एज हैं. i-th डेटा फ़्लो एज, सोर्स x_i, return_value_i और टारगेट y_i, pred_arg_i, body_arg_i के बीच है.

sdy.data_flow_edge, इनपुट के तौर पर किसी एज के रूट टारगेट को लेता है. यह कोई भी टारगेट हो सकता है, लेकिन ब्लॉक आर्ग्युमेंट के बजाय ऑपरेटर का नतीजा होना चाहिए. इसका कोई और इस्तेमाल नहीं होना चाहिए. यह सेशन पूरी तरह से सही नहीं है, क्योंकि यह ऐसा इनपुट ले सकता है जिसका मूल रूप से कोई इस्तेमाल नहीं हुआ है.

sdy.data_flow_edge में, एज के सभी टारगेट के लिए वैकल्पिक तौर पर एक शर्डिंग भी होती है. साथ ही, प्रॉपर्टी के प्रसार के दौरान, टारगेट की शर्डिंग के बजाय उस शर्डिंग को अपडेट किया जाना चाहिए (अगर उसे अटैच किया जा सकता है). यह तब काम आता है, जब किसी ऑप्टिमाइज़ेशन में कई एज होते हैं. ऐसा इसलिए, क्योंकि:

  • हर किनारे से अलग-अलग फ़ॉर्मैट करें.
  • एक साथ सभी टारगेट के बजाय, हर एज की शार्डिंग को अलग-अलग अपडेट करें (उदाहरण के लिए, नतीजों की शार्डिंग के लिए, किसी ऑपरेशन में एक ऐसा TensorShardingPerValueAttr होता है जिसे बदला नहीं जा सकता).
  • जब किसी सोर्स की शार्डिंग बदल गई हो, तो हर किनारे को वर्कलिस्ट में अलग-अलग जोड़ सकता है.

प्रॉपेगेशन, sdy.data_flow_edge के सभी सोर्स और टारगेट के बीच, shardings को इस तरह से प्रॉपेगेट करेगा जैसे कि यह एक सामान्य ऑपरेशन हो. इसमें सोर्स, ऑपरेंड के तौर पर और टारगेट, नतीजों के तौर पर काम करेंगे. साथ ही, एक आइडेंटिटी sdy.op_sharding_rule भी होगी. इसका मतलब है कि फ़ॉरवर्ड प्रोपेगेशन, सोर्स से टारगेट तक होता है और बैकवर्ड प्रोपेगेशन, टारगेट से सोर्स तक होता है.

हम sdy.data_flow_edge के इनपुट को SdyDialect ऑपरेशन से तय करने की अनुमति नहीं देते. इसलिए, हम यह मान सकते हैं कि इसे ऐसे ऑपरेशन से तय किया गया है जिसमें sdy.sharding एट्रिब्यूट रजिस्टर नहीं किया गया है.

खास बातें: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
sharding::mlir::sdy::TensorShardingAttrटेंसर का बंटवारा

ऑपरेंड:

ओपेरैंड ब्यौरा
input किसी भी तरह की वैल्यू के तौर पर

नतीजे:

नतीजा ब्यौरा
result किसी भी तरह की वैल्यू के तौर पर

sdy.manual_computation (sdy::ManualComputationOp)

मैन्युअल कलेक्शन की मदद से, एक से ज़्यादा डिवाइसों पर एक साथ कई काम करना

सिंटैक्स:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

एक्सप्लिसिट कलेक्टिव के साथ हर डिवाइस के लोकल कोड के हिसाब से एक क्षेत्र में जाएं. यहां लॉजिकल शेप हर डिवाइस के फ़िज़िकल बफ़र शेप से मैच करते हैं और कलेक्टिव, फ़िज़िकल क्रॉस-डिवाइस कम्यूनिकेशन से मेल खाते हैं.

बॉडी, manual_axes के हिसाब से स्थानीय है. प्रॉपेगेशन, बॉडी के ज़रिए किसी भी फ़्री ऐक्सिस पर होगा - वे ऐक्सिस जो manual_axes सूची में नहीं हैं.

विशेषताएं: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
in_shardings::mlir::sdy::TensorShardingPerValueAttrकिसी ऑपरेशन के हर ऑपरेंड/नतीजे के हिसाब से टेंसर को अलग-अलग हिस्सों में बांटना
out_shardings::mlir::sdy::TensorShardingPerValueAttrकिसी ऑपरेशन के हर ऑपरेंड/नतीजे के हिसाब से टेंसर को अलग-अलग हिस्सों में बांटना
manual_axes::mlir::sdy::ManualAxesAttr

ऑपरेंड:

ओपेरैंड ब्यौरा
tensors किसी भी तरह की वैल्यू के रैंक किए गए टेंसर का वैरिएडिक

नतीजे:

नतीजा ब्यौरा
results किसी भी तरह के वैल्यू के रैंक किए गए टेंसर का वैरायडिक

sdy.mesh (sdy::MeshOp)

नाम वाला मेश

सिंटैक्स:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

नाम वाले नए मेश को तय करता है. किसी मॉड्यूल में सभी मेश में डिवाइस की संख्या एक जैसी होनी चाहिए (सिर्फ़ एक device_id वाले मेश के अलावा). मेश एक Symbol ऑपरेशन है, जो मॉड्यूल के SymbolTable में दिखता है. साथ ही, इसका रेफ़रंस name से दिया जा सकता है.

विशेषताएं: HasParent<ModuleOp>

इंटरफ़ेस: Symbol

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
sym_name::mlir::StringAttrस्ट्रिंग एट्रिब्यूट
mesh::mlir::sdy::MeshAttrऐक्सिस का मेश और डिवाइसों की सूची

sdy.named_computation (sdy::NamedComputationOp)

नाम वाला कैलकुलेशन ऑपरेशन

सिंटैक्स:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

यह किसी कैलकुलेशन यानी ऑपरेशन के ब्लॉक को ग्रुप में बांटता है और उसे कोई नाम देता है. प्रॉपेगेशन, क्षेत्र में/से उसी तरह से होगा जैसे कि सब कुछ इनलाइन किया गया हो.

इसका इस्तेमाल, कॉल निर्देशों के ज़रिए अन्य फ़ंक्शन में प्रॉपेगेट करने के लिए किया जा सकता है. Shardy के सभी उपयोगकर्ताओं को एक इंपोर्ट/एक्सपोर्ट पास लिखना चाहिए, जो अपने कॉल ऑपरेशन को sdy.named_computation ऑपरेशन में बदल देता है. साथ ही, कॉल किए गए फ़ंक्शन के मुख्य हिस्से को named_computation के मुख्य हिस्से में डुप्लीकेट/कॉपी करता है.

क्षेत्र में मौजूद हर ब्लॉक के आर्ग्युमेंट और रिटर्न की वैल्यू का टाइप, ऑपरेंड के टाइप और ऑपरेशन के नतीजों के टाइप जैसा होना चाहिए.

उदाहरण:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

विशेषताएं: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

इंटरफ़ेस: ConditionallySpeculatable, ShardableDataFlowOpInterface

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
name::mlir::StringAttrस्ट्रिंग एट्रिब्यूट
in_shardings::mlir::sdy::TensorShardingPerValueAttrकिसी ऑपरेशन के हर ऑपरेंड/नतीजे के हिसाब से टेंसर को अलग-अलग हिस्सों में बांटना
out_shardings::mlir::sdy::TensorShardingPerValueAttrकिसी ऑपरेशन के हर ऑपरेंड/नतीजे के हिसाब से टेंसर को अलग-अलग हिस्सों में बांटना

ओपेरैंड:

ओपेरैंड ब्यौरा
operands किसी भी तरह का वैरिएबल

नतीजे:

नतीजा ब्यौरा
«unnamed» किसी भी तरह का वैरिएबल

sdy.propagation_barrier (sdy::PropagationBarrierOp)

प्रोपगेशन बैरियर ऑपरेशन

सिंटैक्स:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

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

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

  • FORWARD का मतलब है कि शार्डिंग सिर्फ़ ऑपरेंड से नतीजे तक जा सकती हैं.
  • BACKWARD का मतलब है कि shardings सिर्फ़ नतीजे से ऑपरेंड में फ़्लो कर सकती हैं.
  • NONE का मतलब है कि इस ऑपरेशन से कोई भी शर्डिंग नहीं हो सकती.
  • BOTH की जानकारी नहीं दी जा सकती, क्योंकि यह ऑपरेशन ज़रूरी नहीं है.

विशेषताएं: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

इंटरफ़ेस: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

इफ़ेक्ट: MemoryEffects::Effect{}

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
allowed_direction::mlir::sdy::PropagationDirectionAttrpropagation direction enum

ऑपरेंड:

ओपेरैंड ब्यौरा
input किसी भी टाइप की वैल्यू के टेंसर की रैंकिंग

नतीजे:

नतीजा ब्यौरा
result किसी भी तरह की वैल्यू का रैंक किया गया टेंसर

sdy.reshard (sdy::ReshardOp)

किसी टेंसर को अलग-अलग शेर्ड में बांटता है

सिंटैक्स:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

इनपुट टेंसर को तय की गई sharding के साथ फिर से शर्ड करता है, जो इनपुट टेंसर की मौजूदा sharding से अलग है.

ShardingConstraintOp और ReshardOp, दोनों ही टेंसर में एक शर्डिंग अटैच करते हैं. इनका लाइफ़स्पैन:

  1. ShardingConstraintOp को उपयोगकर्ता, Sharding propagation से पहले जोड़ते हैं.
  2. शार्डिंग प्रॉपगेशन में ShardingConstraintOp का इस्तेमाल होता है. ShardingConstraintOp, sharding propagation के नतीजों में मौजूद नहीं है. इसके बजाय, ज़रूरत पड़ने पर ReshardOp जोड़ा जा सकता है.
  3. पार्टिशनर, ReshardOp को कलेक्टिव ऑप (या आइडेंटिटी ऑप) में बदल देता है. पार्टीशनर के नतीजों में कोई ReshardOp नहीं होना चाहिए.

// TODO(b/331680067). ग़ैर-ज़रूरी // रीशर्ड ऑपरेशन हटाने के लिए, कैननिकल पैटर्न जोड़ें.

विशेषताएं: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

इंटरफ़ेस: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

इफ़ेक्ट: MemoryEffects::Effect{}

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
sharding::mlir::sdy::TensorShardingAttrटेंसर का बंटवारा

ऑपरेंड:

ओपेरैंड ब्यौरा
input किसी भी टाइप की वैल्यू का टेंसर

नतीजे:

नतीजा ब्यौरा
result किसी भी टाइप की वैल्यू का टेंसर

sdy.return (sdy::ReturnOp)

sdy.return ऑपरेशन, sdy क्षेत्र-आधारित ऑपरेशन से जुड़े इलाकों को खत्म कर देता है. साथ ही, यह शेयर्डी क्षेत्र पर आधारित अन्य ऑपरेशन से जुड़े इलाकों को खत्म कर देता है. यह वैरिएडिक है: यह वैल्यू की एक सूची को आर्ग्युमेंट के तौर पर लेता है. इस सूची में वैल्यू के टाइप अलग-अलग हो सकते हैं, लेकिन एक जैसे होने चाहिए. उदाहरण के लिए, AnyTensor. इसलिए, Shardy IR स्टैक के अलग-अलग लेवल पर इसका फिर से इस्तेमाल किया जा सकता है.

सिंटैक्स:

operation ::= `sdy.return` attr-dict ($results^ `:` type($results))?

खास बातें: AlwaysSpeculatableImplTrait, Terminator

इंटरफ़ेस: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

असर: MemoryEffects::Effect{}

ओपेरैंड:

ओपेरैंड ब्यौरा
results किसी भी तरह की वैरायडिक

sdy.sharding_constraint (sdy::ShardingConstraintOp)

टेंसर को तय शार्डिंग से रोकता है

सिंटैक्स:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

किसी इंटरमीडिएट टेंसर (जैसे, matmul का नतीजा) में एक शर्डिंग अटैच करता है, ताकि यह पता चल सके कि उस टेंसर या उसके इस्तेमाल के सबसेट को इस तरह से शर्ड किया जाना चाहिए.

अगर शीयरिंग में ओपन डाइमेंशन और बिना शर्त वाले ऐक्सिस हैं, तो इसका मतलब है कि टेन्सर को ओपन डाइमेंशन के हिसाब से और भी शीयर किया जा सकता है.

यह ऑपरेशन, इनमें से कोई एक काम कर सकता है:

  • इनका इस्तेमाल न किया गया हो (डैंगलिंग) - इसका मतलब है कि अटैच की गई sharding, इनपुट टेंसर को भी इसी तरह से sharding किया जाना चाहिए.
  • इस्तेमाल किया गया है - इसका मतलब है कि अटैच की गई sharding से पता चलता है कि sharding constraint op के इस्तेमाल को कैसे sharding किया जाना चाहिए. हालांकि, इनपुट टेंसर के अन्य इस्तेमाल के लिए अलग sharding हो सकती है. अगर इनपुट टेंसर का कोई और इस्तेमाल नहीं किया गया है, तो इसका व्यवहार, 'इस्तेमाल नहीं किया गया' केस के जैसा ही होगा.

खास बातें: Elementwise, SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
sharding::mlir::sdy::TensorShardingAttrटेंसर का बंटवारा

ऑपरेंड:

ओपेरैंड ब्यौरा
input किसी भी टाइप की वैल्यू का टेंसर

नतीजे:

नतीजा ब्यौरा
result किसी भी तरह की वैल्यू का टेंसर

sdy.sharding_group (sdy::ShardingGroupOp)

शर्डिंग ग्रुप ऑपरेशन

सिंटैक्स:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

यह ऑप, शार्डिंग ग्रुप ( टेन्सर के ऐसे ग्रुप जिन्हें एक जैसी शार्डिंग रखने के लिए लागू किया जाएगा) को टेंसर असाइन करने के लिए एक इंटरफ़ेस उपलब्ध कराता है. प्रॉपर्टी के डेटा को एक से ज़्यादा शेर्ड में बांटने के दौरान, जैसे ही एक ग्रुप एलिमेंट को शेर्ड किया जाता है, वैसे ही सभी अन्य सदस्यों को भी उसी तरह से शेर्ड कर दिया जाएगा. यह ऑपरेशन, आर्ग्युमेंट ग्रुप आईडी लेता है और कोई नतीजा नहीं दिखाता. इसके बजाय, दिए गए आईडी वाले ग्रुप में इनपुट टेंसर जोड़ने के लिए, इंटरनल sharding ग्रुप के रेप्रज़ेंटेशन में बदलाव करता है.

विशेषताएं:

एट्रिब्यूटMLIR टाइपब्यौरा
group_id::mlir::IntegerAttrबिना साइन वाला 64-बिट इंटिजर एट्रिब्यूट

ऑपरेंड:

ओपेरैंड ब्यौरा
input किसी भी टाइप की वैल्यू के टेंसर की रैंकिंग

विशेषताएं

AxisRefAttr

पूरे ऐक्सिस या स्प्लिट सब-ऐक्सिस का रेफ़रंस

सिंटैक्स:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
नाम ::llvm::StringRef नाम
sub_axis_info SubAxisInfoAttr

DimMappingAttr

किसी डाइमेंशन के लिए फ़ैक्टर इंडेक्स की सूची

सभी फ़ैक्टर इंडेक्स, [0, num_factors) की रेंज में होने चाहिए. खाली सूची से पता चलता है कि यह एक शून्य मैपिंग है. इसे * की मदद से पार्स/प्रिंट किया जाता है. इसका मतलब है कि डाइमेंशन को किसी भी फ़ैक्टर से मैप नहीं किया गया है.

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
factor_indices ::llvm::ArrayRef<int64_t>

DimensionShardingAttr

डाइमेंशन को अलग-अलग हिस्सों में बांटना

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

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
ऐक्सिस ::llvm::ArrayRef<AxisRefAttr> ऐक्सिस रेफ़र की सूची
is_closed bool
प्राथमिकता std::optional<int64_t>

ManualAxesAttr

सिंटैक्स:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
value ::llvm::ArrayRef<StringAttr>

MeshAttr

ऐक्सिस का मेश और डिवाइसों की सूची

सिंटैक्स:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

मेश, ऐक्सिस की सूची होती है. साथ ही, इसमें डिवाइस आईडी की सूची भी हो सकती है. इससे डिवाइस के क्रम की जानकारी मिलती है.

अगर ऐक्सिस की सूची खाली है, तो मेश में साइज़ 1 का एक ऐसा ऐक्सिस होता है जिसका नाम नहीं होता. इस मामले में, अगर डिवाइस आईडी की सूची नहीं दी गई है, तो इंप्लिसिट डिवाइस आईडी सूची [0] होगी. अगर डिवाइस आईडी की सूची दी गई है, तो उसमें किसी भी नॉन-नेगेटिव वैल्यू का एक पूर्णांक शामिल होना चाहिए. हम इसे सबसे ज़्यादा-शर्डिंग वाला मामला कहते हैं.

अगर डिवाइस आईडी की सूची दी गई है, तो ज़्यादा से ज़्यादा शार्डिंग के सभी मामलों में, ऐक्सिस साइज़ का प्रॉडक्ट, डिवाइसों की संख्या से मेल खाना चाहिए. अगर डिवाइस आईडी की सूची नहीं दी गई है, तो डिवाइस आईडी की सूची, iota(product(axes)) होती है. आसानी के लिए, हम ऐसी डिवाइस आईडी सूची तय करने की अनुमति भी नहीं देते जो iota(product(axes)) जैसी हो. इस मामले में, डिवाइस आईडी की सूची तय नहीं की जानी चाहिए.

यहां मेश के कुछ उदाहरण दिए गए हैं:

  • खाली मेश, प्लेसहोल्डर मेश को दिखाता है. इसे प्रॉपेगेशन के दौरान बदला जा सकता है: <[]>
  • बिना नाम वाला ऐक्सिस और साफ़ तौर पर दिखने वाला डिवाइस आईडी वाला मेश. आम तौर पर, इसे ज़्यादा से ज़्यादा शार्डिंग दिखाने के लिए इस्तेमाल किया जाता है: <[], device_ids=[3]>
  • दो ऐक्सिस और डिवाइस आईडी iota(6) के साथ मेश: <["a"=2, "b"=3]>
  • दो ऐक्सिस और डिवाइस के क्रम की जानकारी देने वाले डिवाइस आईडी वाला मेश: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
ऐक्सिस ::llvm::ArrayRef<MeshAxisAttr>
device_ids ::llvm::ArrayRef<int64_t>

MeshAxisAttr

मेश में नाम वाला ऐक्सिस

सिंटैक्स:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
नाम ::llvm::StringRef नाम
साइज़ int64_t

OpShardingRuleAttr

इससे यह तय होता है कि किसी ऑपरेशन को कैसे बांटा जा सकता है.

सिंटैक्स:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  bool   # is_custom_rule
>

शर्डिंग नियम से यह तय होता है कि किसी ऑपरेशन को ऑपरेशन की अलग-अलग प्रॉपर्टी के हिसाब से कैसे बांटा जा सकता है. जैसे, कोई भी एट्रिब्यूट, ऑपरेंड का शेप, नतीजों का शेप वगैरह. उदाहरण के लिए:

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>
%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

ध्यान दें कि हम ऐसे फ़ैक्टर को अनुमति देते हैं जिनका साइज़ 1 है, भले ही उन्हें शेयर नहीं किया जा सकता. ऐसा इसलिए किया जाता है, ताकि सभी ऑपरेशन पूरे किए जा सकें. जैसे, पॉइंटवाइज़ ऑपरेशन में एक डाइमेंशन होते हैं, जो ऑपरेंड और नतीजों से जुड़े होते हैं.

is_custom_rule से पता चलता है कि यह stablehlo.custom_call ऑपरेशन के लिए, उपयोगकर्ता का तय किया गया नियम है या नहीं. पार्टीशन करने वाले टूल को इन ऑपरेशन को पार्टीशन करने का तरीका नहीं पता होता. इसलिए, उपयोगकर्ता को यह बताना होता है. अगर यह कस्टम नियम है, तो नियम को हमेशा सेव रखा जाता है/कभी नहीं हटाया जाता. is_custom_rule सिर्फ़ stablehlo.custom_call ऑपरेशन के लिए सही हो सकता है.

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
factor_sizes ::llvm::ArrayRef<int64_t>
operand_mappings ::llvm::ArrayRef<TensorMappingAttr>
result_mappings ::llvm::ArrayRef<TensorMappingAttr>
is_custom_rule bool

SubAxisInfoAttr

इस सब-ऐक्सिस को पूरे ऐक्सिस से कैसे बनाया गया है, इस बारे में जानकारी

सिंटैक्स:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

किसी पूरे ऐक्सिस को n सब-ऐक्सिस में बांटने पर, ऐक्सिस का आकार [k_1,...,k_n] में बदल जाता है. साथ ही, iवें सब-ऐक्सिस को, ऐक्सिस की बाईं ओर मौजूद सभी साइज़ m=prod(k_1,...,k_(i-1)) (जिसे प्री-साइज़ भी कहा जाता है) और साइज़ k_i के प्रॉडक्ट से दिखाया जा सकता है. इसलिए, सब-ऐक्सिस-जानकारी एट्रिब्यूट में ये दोनों संख्याएं होती हैं और इन्हें इस तरह दिखाया जाता है: प्री-साइज़ m और साइज़ k के लिए (m)k.

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
pre_size int64_t
साइज़ int64_t

TensorMappingAttr

टेंसर के हर डाइमेंशन के लिए फ़ैक्टर मैपिंग.

सिंटैक्स:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
dim_mappings ::llvm::ArrayRef<DimMappingAttr>

TensorShardingAttr

टेंसर का बंटवारा

सिंटैक्स:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>   # replicated_axes
>

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

इस शर्डिंग के लिए टारगेट किए गए मेश को, सिंबल के नाम से दिखाया जा सकता है. इसके लिए, उससे जुड़े MeshOp सिंबल या इनलाइन किए गए MeshAttr का रेफ़रंस दिया जा सकता है.

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
mesh_or_ref ::mlir::Attribute mesh attr या flat mesh symbol reference attr
dim_shardings ::llvm::ArrayRef<DimensionShardingAttr>
replicated_axes ::llvm::ArrayRef<AxisRefAttr> ऐक्सिस रेफ़रंस की सूची

TensorShardingPerValueAttr

किसी ऑपरेशन के हर ऑपरेंड/नतीजे के हिसाब से टेंसर को अलग-अलग हिस्सों में बांटना

सिंटैक्स:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

पैरामीटर:

पैरामीटर C++ टाइप ब्यौरा
shardings ::llvm::ArrayRef<TensorShardingAttr>

Enums

PropagationDirection

propagation direction enum

केस:

चिह्न मान स्ट्रिंग
कोई नहीं 0 कोई नहीं
FORWARD 1 FORWARD
BACKWARD 2 BACKWARD
दोनों 3 दोनों