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::PropagationDirectionAttr | propagation direction enum |
ऑपरेंड:
ओपेरैंड | ब्यौरा |
---|---|
input |
किसी भी टाइप की वैल्यू के टेंसर की रैंकिंग |
नतीजे:
नतीजा | ब्यौरा |
---|---|
result |
किसी भी तरह की वैल्यू का रैंक किया गया टेंसर |
sdy.reshard
(sdy::ReshardOp)
किसी टेंसर को अलग-अलग शेर्ड में बांटता है
सिंटैक्स:
operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)
इनपुट टेंसर को तय की गई sharding के साथ फिर से शर्ड करता है, जो इनपुट टेंसर की मौजूदा sharding से अलग है.
ShardingConstraintOp और ReshardOp, दोनों ही टेंसर में एक शर्डिंग अटैच करते हैं. इनका लाइफ़स्पैन:
- ShardingConstraintOp को उपयोगकर्ता, Sharding propagation से पहले जोड़ते हैं.
- शार्डिंग प्रॉपगेशन में ShardingConstraintOp का इस्तेमाल होता है. ShardingConstraintOp, sharding propagation के नतीजों में मौजूद नहीं है. इसके बजाय, ज़रूरत पड़ने पर ReshardOp जोड़ा जा सकता है.
- पार्टिशनर, 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 |
दोनों |