'sdy' बोली

शार्डी (एसडीवाय) बोली

शार्डी (SDY) डायलेक्ट, ऐक्सिस पर आधारित टेंसर शार्डिंग और टेंसर में शार्डिंग जोड़ने के लिए अतिरिक्त एपीआई कॉम्पोनेंट तय करता है.

वर्शन लॉग: 0.0.1: TensorShardingAttr में, कम नहीं किए गए ऐक्सिस जोड़ें.

कार्रवाइयां

sdy.all_gather (sdy::AllGatherOp)

एक्सिस के हिसाब से, सभी डिवाइसों के बीच डेटा शेयर करता है

सिंटैक्स:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

यह फ़ंक्शन, gathering_axes में तय किए गए ऐक्सिस के हिसाब से, किसी टेंसर के हिस्सों को इकट्ठा करता है.

gathering_axes, ऐक्सिस की सूचियों की सूची होती है. बाहरी सूची, टेंसर के डाइमेंशन से ज़्यादा है. हर इनर लिस्ट में, उन ऐक्सिस के बारे में बताया जाता है जिनके हिसाब से, हर डाइमेंशन पर अलग-अलग तरीके से डेटा इकट्ठा किया जाना चाहिए. इसे ऑपरेंड (tensor) की शार्डिंग पर लागू किया जाएगा, ताकि नतीजे (out_sharding) की शार्डिंग मिल सके.

ध्यान दें कि out_sharding का इस्तेमाल, नतीजे को अलग-अलग हिस्सों में बांटने के लिए नहीं किया जाता. इसके बजाय, नतीजे की शार्डिंग, ऑपरेंड और gathering_axes की शार्डिंग से तय होती है. साथ ही, out_sharding को इस अनुमानित शार्डिंग से मेल खाना चाहिए.

उदाहरण:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• gathering_axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• ऑपरेंड शार्डिंग पर gathering_axes लागू करने से out_sharding मिलता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.all_reduce (sdy::AllReduceOp)

ऐक्सिस के हिसाब से, सभी कम्यूनिकेशन को कम करना

सिंटैक्स:

operation ::= `sdy.all_reduce` ($reduction_op^)? $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

यह फ़ंक्शन, reduction_axes में तय किए गए ऐक्सिस के हिसाब से, किसी टेंसर के हिस्सों को कम करता है. नतीजे के लिए, reduction_axes का क्रम ज़रूरी नहीं है. हालांकि, इससे मिलते-जुलते ग्रुप के क्रम पर असर पड़ सकता है.

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• reduction_axes को AxisRefListAttr में दी गई शर्तों को पूरा करना होगा.
• reduction_axes को मेश के हिसाब से क्रम से लगाया जाना चाहिए.
• ऑपरेंड शार्डिंग और out_sharding में डाइमेंशन शार्डिंग एक जैसी होनी चाहिए.
• reduction_axes, ऑपरेंड डाइमेंशन शार्डिंग और रेप्लिकेटेड ऐक्सिस के साथ ओवरलैप नहीं होना चाहिए. हालांकि, यह अनरिड्यूस्ड ऐक्सिस के साथ ओवरलैप हो सकता है.
• reduction_axes को out_sharding के बिना कम किए गए ऐक्सिस के साथ ओवरलैप नहीं होना चाहिए. दूसरे शब्दों में, out_sharding को reduction_axes के साथ दोहराया जाना चाहिए. ऐसा साफ़ तौर पर या परोक्ष रूप से किया जा सकता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.all_slice (sdy::AllSliceOp)

यह फ़ंक्शन, ऐक्सिस के साथ डाइनैमिक-स्लाइस ऑपरेशन करता है

सिंटैक्स:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

यह फ़ंक्शन, slicing_axes में बताए गए ऐक्सिस के हिसाब से, किसी टेंसर के हिस्सों को स्लाइस करता है. sdy.all_slice और sdy.all_gather के बीच बीजगणितीय द्वैतता है.

slicing_axes, ऐक्सिस की सूचियों की सूची होती है. बाहरी सूची, टेंसर के डाइमेंशन से ज़्यादा है. हर इनर लिस्ट में, उन ऐक्सिस के बारे में बताया जाता है जिनके हिसाब से डाइमेंशन को स्लाइस किया जाना चाहिए. इसे ऑपरेंड (tensor) के शार्डिंग पर लागू किया जाएगा, ताकि नतीजे (out_sharding) का शार्डिंग मिल सके.

ध्यान दें कि out_sharding का इस्तेमाल, नतीजे को अलग-अलग हिस्सों में बांटने के लिए नहीं किया जाता. इसके बजाय, नतीजे की शार्डिंग, ऑपरेंड और slicing_axes की शार्डिंग से तय होती है. साथ ही, out_sharding को इस अनुमानित शार्डिंग से मेल खाना चाहिए.

उदाहरण:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• slicing_axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• ऑपरेंड शार्डिंग पर slicing_axes लागू करने से out_sharding मिलता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.all_to_all (sdy::AllToAllOp)

एक्सिस के हिसाब से सभी के साथ कम्यूनिकेट करता है

सिंटैक्स:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

पैरामीटर की सूची में मौजूद हर (axes, src_dim, tgt_dim) टपल के लिए, यह ऑपरेशन, tgt_dim डाइमेंशन और axes में बताए गए ऐक्सिस के हिसाब से, टेंसर के हिस्सों को स्लाइस करता है. साथ ही, उन हिस्सों को ऐक्सिस के हिसाब से स्कैटर करता है और उन्हें src_dim डाइमेंशन के हिसाब से जोड़ता है.

यह कार्रवाई, src_dim और axes के साथ-साथ ऑल-गैदर और tgt_dim और axes के साथ-साथ ऑल-स्लाइस का कॉम्बिनेशन है. इसका मतलब है कि इनपुट टेंसर पर मौजूद ऐक्सिस शार्डिंग डाइमेंशन src_dim का सफ़िक्स, आउटपुट टेंसर पर मौजूद ऐक्सिस शार्डिंग डाइमेंशन tgt_dim में जोड़ दिया जाता है.

ऑपरेंड (tensor) की शार्डिंग पर ऑल-टू-ऑल लागू किया जाएगा, ताकि नतीजे (out_sharding) की शार्डिंग मिल सके.

ध्यान दें कि out_sharding का इस्तेमाल, नतीजे को अलग-अलग हिस्सों में बांटने के लिए नहीं किया जाता. इसके बजाय, नतीजे की शार्डिंग, ऑपरेंड src_dim, tgt_dim, और axes की शार्डिंग से तय होती है. साथ ही, out_sharding को इस अनुमानित शार्डिंग से मेल खाना चाहिए.

उदाहरण:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• पैरामीटर की सूची वाला फ़ील्ड खाली नहीं छोड़ा जा सकता.
• params में मौजूद हर पैरामीटर के लिए:
  • axes में मौजूद एलिमेंट, AxisRefAttr की शर्तों को पूरा करते हों.
  • src_dim और tgt_dim, मान्य डाइमेंशन होने चाहिए. ये नॉन-नेगेटिव होने चाहिए और इनकी वैल्यू, टेंसर के रैंक से कम होनी चाहिए.
  • सभी पैरामीटर में src_dim या tgt_dim यूनीक होना चाहिए.
  • src_dim को सभी पैरामीटर के हिसाब से बढ़ते क्रम में लगाया जाना चाहिए.
• ऑपरेंड शार्डिंग में, src_dim से tgt_dim में axes को ले जाने पर out_sharding मिलता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.collective_permute (sdy::CollectivePermuteOp)

ऐक्सिस बदलने के लिए, कलेक्टिव-परम्यूट कम्यूनिकेशन करता है

सिंटैक्स:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

यह फ़ंक्शन, हर डिवाइस से इनपुट टेंसर का एक हिस्सा दूसरे डिवाइस को भेजता है, ताकि टेंसर को शार्ड करने वाले ऐक्सिस को फिर से क्रम में लगाया जा सके या बदला जा सके.

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

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

यहां दिए गए उदाहरण में, शार्ड किए गए टेंसर का साइज़ tensor<1x4x2xf32> है. इसे कलेक्टिव परम्यूटेशन के ज़रिए सुरक्षित रखा जाता है.

उदाहरण:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• अगर इनपुट और आउटपुट शार्डिंग में अलग-अलग मेश हैं, तो उन मेश में एक जैसे ऐक्सिस और डिवाइस आईडी का क्रम अलग-अलग होना चाहिए.
• हर डाइमेंशन के लिए, out_sharding में शार्डिंग ऐक्सिस के साइज़ का प्रॉडक्ट, ऑपरेंड डाइमेंशन की शार्डिंग से मेल खाना चाहिए.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

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{}

विशेषताएं:

नतीजे:

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 लूप में 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 के सभी सोर्स और टारगेट के बीच शार्डिंग को इस तरह से फैलाएगा जैसे कि यह सोर्स को ऑपरेंड और टारगेट को नतीजे के तौर पर इस्तेमाल करने वाला एक सामान्य ऑपरेशन हो. साथ ही, इसमें एक आइडेंटिटी sdy.op_sharding_rule होगी. इसका मतलब है कि फ़ॉरवर्ड प्रोपगेशन, सोर्स से टारगेट तक होता है और बैकवर्ड प्रोपगेशन, टारगेट से सोर्स तक होता है.

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

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.func_data_flow_edge (sdy::FuncDataFlowEdgeOp)

फ़ंक्शन के इनपुट/आउटपुट डेटा फ़्लो का एज ऑप.

सिंटैक्स:

operation ::= `sdy.func_data_flow_edge` $operand attr-dict `:` type($result)

यह फ़ंक्शन के आर्ग्युमेंट या कॉल के नतीजों के लिए, डेटा फ़्लो एज ऑप है. जब इसका ऑपरेंड BlockArgument होता है, तब यह caller callOp के आर्ग्युमेंट से लेकर func आर्ग्युमेंट के उपयोगकर्ताओं तक एक ब्रिज होता है. हर फ़ंक्शन आर्ग्युमेंट के लिए, एक फ़ंक्शन डेटा फ़्लो एज होता है. जब इसका ऑपरेंड OpResult होता है, तो यह कॉल किए गए funcOp की रिटर्न वैल्यू से कॉल के नतीजे के उपयोगकर्ताओं तक एक ब्रिज होता है. हर कॉल के नतीजे के लिए, एक फ़ंक्शन डेटा फ़्लो एज होता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, SymbolUserOpInterface

ऑपरेंड:

नतीजे:

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 की सूची में शामिल नहीं होते.

ध्यान दें कि रैंक न किए गए किसी भी टेंसर में, रैंक 0 के साथ शार्डिंग होनी चाहिए. इसका मतलब है कि उसे पूरी तरह से दोहराया गया हो.

चुनौतियां:

• in_shardings और out_shardings में मौजूद एलिमेंट को, TensorShardingAttr में दी गई पाबंदियों का पालन करना होगा.
• ऑपरेशन क्षेत्र के ग्लोबल और लोकल टेंसर इनपुट/आउटपुट की संख्या मेल खानी चाहिए.
• हर डिम शार्डिंग में, मैन्युअल ऐक्सिस, फ़्री ऐक्सिस से पहले होने चाहिए.
• मैन्युअल ऐक्सिस में पैडिंग नहीं जोड़ी जा सकती. इसका मतलब है कि डाइमेंशन का साइज़, मैन्युअल ऐक्सिस के साइज़ से पूरी तरह से भाग होना चाहिए.
• ओपी रीज़न के आर्ग्युमेंट/नतीजों के ग्लोबल और लोकल शेप मेल खाने चाहिए.

ट्रेट: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

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

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.mesh (sdy::MeshOp)

नाम वाला मेश

सिंटैक्स:

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

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

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

इंटरफ़ेस: Symbol

विशेषताएं:

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 के मुख्य हिस्से में डुप्लीकेट/कॉपी किया जा सकता है.

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

उदाहरण:

%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, InferTypeOpInterface, ShardableDataFlowOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

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

सिंटैक्स:

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

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

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

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

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

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

असर: MemoryEffects::Effect{}

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.reduce_scatter (sdy::ReduceScatterOp)

यह फ़ंक्शन, ऐक्सिस के साथ-साथ कम्यूनिकेशन को कम करता है

सिंटैक्स:

operation ::= `sdy.reduce_scatter` ($reduction_op^)? $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

यह फ़ंक्शन, reduce_scatter_axes में दिए गए ऐक्सिस के हिसाब से, किसी टेंसर के चंक को कम करता है. इसके बाद, नतीजे को उन्हीं ऐक्सिस के हिसाब से स्कैटर करता है. यह ऑपरेशन, reduce_scatter_axes के साथ-साथ sdy.all_reduce और sdy.all_slice का कॉम्बिनेशन होता है.

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• reduce_scatter_axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• ऑपरेंड शार्डिंग पर reduce_scatter_axes लागू करने से out_sharding मिलता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: CollectiveOpInterface, InferTypeOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

जिन ऐक्सिस को कम नहीं किया गया है उनमें, साफ़ तौर पर या अपने-आप दोहराए गए ऐक्सिस को ले जाएं.

सिंटैक्स:

operation ::= `sdy.replicated_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

ऑपरेंड में axes को साफ़ तौर पर या परोक्ष रूप से दोहराया जाना चाहिए. इस कार्रवाई से, खोज के नतीजों में उनकी रैंकिंग कम नहीं होती. हमारा संबंध इस तरह का है:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

उदाहरण:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}, {}\], replicated={"c", "d"}, unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.replicated_to_unreduced {"a", "c", "f"} %1 out_sharding=<@mesh, [{"b"}, {}, {}\], replicated={"d"}, unreduced={"a", "c", "e", "f"}> : tensor<8x8x8xf32>

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• axes को AxisRefListAttr में दी गई शर्तों को पूरा करना होगा.
• axes को मेश के हिसाब से क्रम से लगाया जाना चाहिए.
• axes खाली नहीं हैं.
• इनपुट और आउटपुट शार्डिंग में डाइमेंशन शार्डिंग एक जैसी होनी चाहिए.
• axes को ऑपरेंड शार्डिंग में, परोक्ष या प्रत्यक्ष तौर पर दोहराया जाना चाहिए.
• inUnreducedAxes + axes = outUnreducedAxes.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.reshard (sdy::ReshardOp)

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

सिंटैक्स:

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

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

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

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

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

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

असर: MemoryEffects::Effect{}

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.return (sdy::ReturnOp)

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

सिंटैक्स:

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

विशेषताएं: AlwaysSpeculatableImplTrait, ReturnLike, Terminator

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

असर: MemoryEffects::Effect{}

ऑपरेंड:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

ऑपरेंड के कुछ शार्ड किए गए ऐक्सिस को, नतीजे के अनरिड्यूस्ड ऐक्सिस में ले जाएं.

सिंटैक्स:

operation ::= `sdy.sharded_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

ऑपरेंड को शार्ड करने के लिए, axes का इस्तेमाल किया जाना चाहिए. इस कार्रवाई से, नतीजे में उनकी संख्या कम नहीं होती. हमारा संबंध इस तरह का है:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), where all-gather, sharded-to-unreduced, all-reduce are applied on the same axes.

उदाहरण:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\], unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.sharded_to_unreduced [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\], unreduced={"b", "c", "d", "e"}> : tensor<8x8x8xf32>

चुनौतियां:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना ज़रूरी है.
• axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• ऑपरेंड शार्डिंग पर axes लागू करने से out_sharding मिलता है.

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, Sdy_CollectiveOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

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

सिंटैक्स:

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

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

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

यह कार्रवाई इनमें से कोई एक काम कर सकती है:

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

विशेषताएं: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface, SymbolUserOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.sharding_group (sdy::ShardingGroupOp)

ग्रुप में मौजूद टेंसर को एक ही शार्डिंग में सीमित करता है.

सिंटैक्स:

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

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

इंटरफ़ेस: InferTypeOpInterface

विशेषताएं:

ऑपरेंड:

विशेषताएं

AllToAllParamAttr

ऑल-टू-ऑल पैरामीटर

सिंटैक्स:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

यह एक टपल है. इसमें सभी-से-सभी ऑपरेशन करने के लिए, ऐक्सिस और सोर्स/टारगेट डाइमेंशन शामिल होते हैं.

पैरामीटर:

AllToAllParamListAttr

सभी-से-सभी पैरामीटर की सूची

सिंटैक्स:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

पैरामीटर:

AxisRefAttr

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

सिंटैक्स:

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

चुनौतियां:

• name, MeshAttr में मौजूद होना चाहिए.
• अगर sub_axis_info मौजूद है, तो यह ज़रूरी है कि यह SubAxisInfoAttr की शर्तों को पूरा करता हो.

पैरामीटर:

AxisRefListAttr

ऐक्सिस रेफ़रंस की सूची

सिंटैक्स:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

चुनौतियां:

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

पैरामीटर:

AxisToPropagationDetailsAttr

किसी ऐक्सिस और सोर्स के लिए, एज फ़्लो की जानकारी.

सिंटैक्स:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

यह फ़ंक्शन, किसी सोर्स वैल्यू रेफ़रंस को किसी खास ऐक्सिस के साथ टारगेट वैल्यू रेफ़रंस की सूची पर मैप करता है.

पैरामीटर:

DimMappingAttr

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

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

चुनौतियां:

• कम से कम एक फ़ैक्टर इंडेक्स मौजूद है.
• फ़ैक्टर इंडेक्स, [0, $factor_sizes) की रेंज में होने चाहिए.
• अगर एक से ज़्यादा फ़ैक्टर हैं, तो उनमें से किसी का भी साइज़ 1 नहीं हो सकता.
• डुप्लीकेट फ़ैक्टर इंडेक्स नहीं होने चाहिए.

पैरामीटर:

DimensionShardingAttr

डाइमेंशन शार्डिंग

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

चुनौतियां:

• axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• अगर डाइमेंशन शार्डिंग की कोई प्राथमिकता है, तो:
  • प्राथमिकता 0 या इससे ज़्यादा है.
  • अगर डाइमेंशन बंद है, तो उसमें कम से कम एक ऐक्सिस होता है.

पैरामीटर:

EdgeValueRefAttr

type टाइप के वैल्यू एज के किसी इंडेक्स का रेफ़रंस.

सिंटैक्स:

#sdy.edge_value_ref<
  `operand` | `result`,   # type
  int64_t   # index
>

पैरामीटर:

ListOfAxisRefListsAttr

ऐक्सिस रेफ़रंस लिस्ट की सूची

सिंटैक्स:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

पैरामीटर:

ManualAxesAttr

उन ऐक्सिस की सूची जिन पर ManualComputationOp मैन्युअल है

सिंटैक्स:

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

पैरामीटर:

MeshAttr

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

सिंटैक्स:

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

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

अगर ऐक्सिस की सूची खाली है

• अगर device_ids नहीं दिया गया है, तो यह एक खाली मेश है.
• अगर device_ids दिया गया है, तो यह एक नॉन-नेगेटिव पूर्णांक होना चाहिए. इसे मैक्सिमल-शार्डिंग मेश कहा जाता है.

अगर ऐक्सिस की सूची दी गई है

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

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

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

चुनौतियां:

• device_ids में मौजूद एलिमेंट, नेगेटिव नहीं होने चाहिए.
• अगर axes खाली है, तो device_ids का साइज़ 0 (खाली मेश) या 1 (ज़्यादा से ज़्यादा शार्डिंग वाला मेश) हो सकता है.
• अगर axes खाली नहीं है, तो
  • axes में मौजूद एलिमेंट के नाम डुप्लीकेट नहीं होने चाहिए.
  • अगर device_ids को तय किया जाता है, तो मूल device_ids iota(product(axis_sizes)) नहीं होता और क्रम से लगाया गया device_ids iota(product(axis_sizes)) होता है.

पैरामीटर:

MeshAxisAttr

मेश में नाम दिया गया ऐक्सिस

सिंटैक्स:

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

पैरामीटर:

OpShardingRuleAttr

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

सिंटैक्स:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  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 वाले फ़ैक्टर की अनुमति देते हैं, भले ही उन्हें शार्ड न किया जा सके. ऐसा मुख्य रूप से इसलिए किया जाता है, ताकि जानकारी पूरी हो. ऐसा इसलिए, क्योंकि पॉइंटवाइज़ ऑप्स जैसे कई ऑप्स में साइज़ एक डाइमेंशन होता है, जो ऑपरेंड और नतीजों के हिसाब से होता है.

फ़ैक्टर टाइप:

• reduction_factors में उन फ़ैक्टर के इंडेक्स होते हैं जिन्हें कम करने की ज़रूरत होती है. जैसे, डॉट ऑपरेशन में कॉन्ट्रैक्टिंग डाइमेंशन. ये फ़ैक्टर, ऑपरेंड में हो सकते हैं, लेकिन नतीजों में नहीं.
• need_replication_factors में उन फ़ैक्टर के इंडेक्स शामिल होते हैं जिनके लिए पूरी तरह से रेप्लिकेशन की ज़रूरत होती है. जैसे, सॉर्ट करने की कार्रवाई में सॉर्ट किया गया डाइमेंशन.
• permutation_factors में उन फ़ैक्टर के इंडेक्स होते हैं जिनके लिए collective-permute की ज़रूरत होती है. ऐसा तब होता है, जब उन्हें शार्ड किया जाता है. जैसे, पैड ऑपरेशन में पैडिंग डाइमेंशन.
• अन्य सभी फ़ैक्टर को पास-थ्रू फ़ैक्टर माना जाता है. इसका मतलब है कि अगर सभी टेंसर में एक ही तरीके से शार्ड किया जाता है, तो इन फ़ैक्टर के लिए किसी भी तरह के कम्यूनिकेशन की ज़रूरत नहीं होती.

blocked_propagation_factors में वे फ़ैक्टर शामिल होते हैं जिनके हिसाब से शार्डिंग को आगे बढ़ाने की अनुमति नहीं होती. यह फ़ैक्टर टाइप के ऑर्थोगोनल है. जैसे, ब्लॉक किए गए प्रोपगेशन फ़ैक्टर, किसी भी तरह के फ़ैक्टर हो सकते हैं.

is_custom_rule से पता चलता है कि यह उपयोगकर्ता की ओर से तय किया गया नियम है या नहीं. उपयोगकर्ता, अपने कस्टम कॉल के लिए शार्डिंग के नियम तय कर सकते हैं. इसके अलावा, वे स्टैंडर्ड ऑपरेशन के लिए पहले से तय किए गए शार्डिंग के नियमों को बदल सकते हैं. कस्टम नियम को हमेशा सेव किया जाता है/कभी नहीं हटाया जाता.

चुनौतियां:

• ऑपरेंड/नतीजे के मैपिंग की संख्या, ऑप के ऑपरेंड/नतीजों की संख्या से मेल खानी चाहिए.
• कम से कम एक मैपिंग मौजूद है. ऐसा नहीं हो सकता कि किसी ऐसे ऑपरेटर के लिए नियम बनाया गया हो जिसके कोई ऑपरेंड/नतीजे न हों.
• हर TensorMappingAttr की रैंक, उससे जुड़े टेंसर टाइप की रैंक से मेल खाती है.
• फ़ैक्टर के हर ग्रुप (reduction_factors, need_replication_factors, permutation_factors) के लिए:
  • एलिमेंट, [0, $factor_sizes] की सीमा में होने चाहिए.
  • हर ग्रुप और सभी ग्रुप में, फ़ैक्टर इंडेक्स डुप्लीकेट नहीं होने चाहिए.

पैरामीटर:

PropagationEdgesAttr

प्रॉपगेशन के सभी चरणों के लिए, एज मेटाडेटा का प्रॉपगेशन.

सिंटैक्स:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

किसी वैल्यू के लिए, हर ऐक्सिस के हिसाब से आगे बढ़ने की जानकारी की सूची. इसे चरण के इंडेक्स के हिसाब से ग्रुप किया जाता है.

पैरामीटर:

PropagationOneStepAttr

हर चरण के लिए, मेटाडेटा को आगे बढ़ाने की जानकारी.

सिंटैक्स:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

प्रोपैगेशन के एक चरण के लिए, सभी ऐक्सिस के लिए प्रोपैगेशन की जानकारी.

पैरामीटर:

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 के प्रॉडक्ट के तौर पर दिखाया जा सकता है. इसलिए, sub-axis-info एट्रिब्यूट में ये दो नंबर होते हैं और इसे इस तरह दिखाया जाता है: प्री-साइज़ m और साइज़ k के लिए (m)k.

चुनौतियां:

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

पैरामीटर:

TensorMappingAttr

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

सिंटैक्स:

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

चुनौतियां:

• dim_mappings में मौजूद एलिमेंट, DimMappingAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• किसी भी डाइमेंशन में, डुप्लीकेट फ़ैक्टर इंडेक्स नहीं होने चाहिए.

पैरामीटर:

TensorShardingAttr

टेंसर शार्डिंग

सिंटैक्स:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>,   # unreduced_axes
  `sum` | `max` | `min`   # reduction_op
>

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

ध्यान दें कि किसी टेंसर पर कोई शार्डिंग एट्रिब्यूट न होने का मतलब है कि टेंसर शार्डिंग पूरी तरह से खुली है.

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

शार्डिंग में ऐसे ऐक्सिस हो सकते हैं जिनमें कमी नहीं की गई है. इन्हें unreduced_axes से तय किया जाता है. इसका मतलब है कि इन ऐक्सिस के हिसाब से, टेंसर में कमी नहीं की गई है. उदाहरण के लिए, अगर matmul का कॉन्ट्रैक्टिंग डाइमेंशन, lhs और rhs दोनों में x ऐक्सिस के साथ शार्ड किया गया है, तो नतीजे को x के साथ कम नहीं किया जाता. अनरिड्यूस्ड ऐक्सिस के साथ-साथ, टेंसर पर all-reduce लागू करने से, टेंसर उन ऐक्सिस के साथ रेप्लिकेट हो जाएगा. हालांकि, कम किए गए ऐक्सिस वाले टेंसर को तुरंत पूरी तरह से कम करने की ज़रूरत नहीं होती. इसे stablehlo.add जैसे लीनियर ऑपरेशन में पास करने पर, कम नहीं किया जा सकता. हालांकि, बाद में इसे पूरी तरह से कम किया जा सकता है. हम मानते हैं कि कटौती का टाइप sum है. आने वाले समय में, अन्य कटौतियों के लिए सहायता उपलब्ध कराई जा सकती है.

चुनौतियां:

• dim_shardings में मौजूद एलिमेंट, DimensionShardingAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• replicated_axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• unreduced_axes में मौजूद एलिमेंट, AxisRefListAttr में दी गई शर्तों के मुताबिक होने चाहिए.
• अगर इससे जुड़ा टेंसर टाइप ShapedType नहीं है, तो शार्डिंग की रैंक 0 होनी चाहिए और इसमें कोई भी दोहराया गया ऐक्सिस नहीं होना चाहिए.
• अगर यह ShapedType है, तो:
  • टेंसर की कोई रैंक होनी चाहिए.
  • डाइमेंशन शार्डिंग की संख्या, टेंसर की रैंक के बराबर होती है.
  • साइज़ 0 के डाइमेंशन को शार्ड नहीं किया जाता.
• dim_shardings, replicated_axes, और unreduced_axes में, एक-दूसरे से ओवरलैप होने वाले डुप्लीकेट ऐक्सिस-रेफ़रंस या सब-ऐक्सिस मौजूद नहीं हैं.
• replicated_axes और unreduced_axes में मौजूद आइटम, mesh_or_ref के हिसाब से क्रम में लगाए जाते हैं (AxisRefAttr::getMeshComparator देखें).

पैरामीटर:

TensorShardingPerValueAttr

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

सिंटैक्स:

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

TensorShardingAttr की सूची. इसमें हर ऑपरेंड/ऑपरेशन के नतीजे के लिए एक होता है.

चुनौतियां:

• shardings में मौजूद एलिमेंट, TensorShardingAttr की शर्तों को पूरा करते हों.

पैरामीटर:

Enums

EdgeNodeType

Edge node type enum

केस:

PropagationDirection

Propagation direction enum

केस:

ReductionOp

Reduction op enum

केस: