'sdy' बोली

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

ऑपरेशंस

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

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.all_reduce (sdy::AllReduceOp)

ऐक्सिस के साथ-साथ सभी एलिमेंट के लिए कम्यूनिकेशन करना

सिंटैक्स:

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

reduction_axes में बताए गए ऐक्सिस के हिसाब से, टेंसर के चंक कम कर देता है. नतीजे के लिए reduction_axes का क्रम ज़रूरी नहीं है. हालांकि, इससे उससे जुड़े डुप्लीकेट ग्रुप के क्रम पर असर पड़ सकता है.

सीमाएं:

• Sdy_CollectiveOpInterface में दी गई शर्तों को पूरा करना चाहिए.
• reduction_axes को AxisRefListAttr में दी गई शर्तों को पूरा करना होगा;
• reduction_axes, ऑपरेंड के लिए इस्तेमाल होने वाले, डेटा को अलग-अलग हिस्सों में बांटने वाले ऐक्सिस से ओवरलैप नहीं होना चाहिए;

खास बातें: SameOperandsAndResultType

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

विशेषताएं:

ऑपरेंड:

नतीजे:

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, ऐक्सिस की सूचियों की सूची है. बाहरी सूची, टेंसर के डाइमेंशन से ज़्यादा है. हर इनर सूची उन ऐक्सिस के बारे में बताती है जिनके हिसाब से, संबंधित डाइमेंशन पर स्लाइस की जानी चाहिए. यह नतीजे (out_sharding) की शर्डिंग पाने के लिए, ऑपरेंड (tensor) की शर्डिंग पर लागू किया जाएगा.

ध्यान दें कि नतीजे को अलग-अलग हिस्सों में बांटने के लिए, 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>

सीमाएं:

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

खास बातें: SameOperandsAndResultType

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

विशेषताएं:

ऑपरेंड:

नतीजे:

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 में जोड़ा जाता है.

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

ध्यान दें कि नतीजे को अलग-अलग हिस्सों में बांटने के लिए, 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 को सभी पैरामीटर में बढ़ते क्रम में क्रम से लगाया जाना चाहिए.
• ऑपरेंड को अलग-अलग हिस्सों में बांटने की सुविधा में, axes को src_dim से tgt_dim में ले जाने पर, out_sharding मिलता है.

खास बातें: SameOperandsAndResultType

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

विशेषताएं:

ऑपरेंड:

नतीजे:

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 में दी गई शर्तों को पूरा करना चाहिए.
• अगर इनपुट और आउटपुट sharding में अलग-अलग मेश हैं, तो उन मेश में एक जैसे ऐक्सिस होने चाहिए और डिवाइस आईडी का क्रम अलग-अलग होना चाहिए.
• हर डाइमेंशन के लिए, out_sharding में मौजूद, 'शर्डिंग ऐक्सिस के साइज़' का प्रॉडक्ट, ऑपरेंड डाइमेंशन की शर्डिंग से मैच करना चाहिए.

खास बातें: SameOperandsAndResultType

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

विशेषताएं:

ऑपरेंड:

नतीजे:

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

विशेषताएं:

ऑपरेंड:

नतीजे:

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 सूची में नहीं हैं.

सीमाएं:

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

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

इंटरफ़ेस: ShardableDataFlowOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.mesh (sdy::MeshOp)

नाम वाला मेश

सिंटैक्स:

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

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

खास बातें: HasParent<ModuleOp>

इंटरफ़ेस: 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 के मुख्य हिस्से में जोड़ना होगा.

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

उदाहरण:

%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

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

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

सिंटैक्स:

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

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

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

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

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

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

असर: MemoryEffects::Effect{}

विशेषताएं:

ऑपरेंड:

नतीजे:

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, SameOperandsAndResultType

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

असर: MemoryEffects::Effect{}

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.return (sdy::ReturnOp)

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

सिंटैक्स:

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

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

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

असर: MemoryEffects::Effect{}

ऑपरेंड:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

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

सिंटैक्स:

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

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

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

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

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

खास बातें: SameOperandsAndResultType

इंटरफ़ेस: InferTypeOpInterface

विशेषताएं:

ऑपरेंड:

नतीजे:

sdy.sharding_group (sdy::ShardingGroupOp)

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

सिंटैक्स:

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

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

इंटरफ़ेस: 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 की शर्तों को पूरा करने चाहिए.
• डुप्लीकेट ऐक्सिस-रेफ़रंस या सब-ऐक्सिस मौजूद न हों, जो एक-दूसरे से ओवरलैप करते हों.
• आस-पास मौजूद कोई भी दो ऐक्सिस-रेफ़रंस, एक ही फ़ुल ऐक्सिस के लगातार सब-ऐक्सिस नहीं हैं. इसका मतलब है कि उन्हें एक सब-ऐक्सिस या फ़ुल ऐक्सिस में मर्ज किया जा सकता है.

पैरामीटर:

DimMappingAttr

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

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

सीमाएं:

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

पैरामीटर:

DimensionShardingAttr

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

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

सीमाएं:

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

पैरामीटर:

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
>

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

अगर ऐक्सिस की सूची खाली है, तो मेश में साइज़ 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]>

सीमाएं:

• axes में मौजूद एलिमेंट के नाम डुप्लीकेट नहीं होने चाहिए.
• अगर device_ids तय किया गया है, तो:
  • ऐक्सिस के साइज़ का प्रॉडक्ट, डिवाइसों की संख्या से मेल खाना चाहिए.
  • इसके सभी एलिमेंट, 0 से ज़्यादा होने चाहिए.
  • 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 में उन फ़ैक्टर के इंडेक्स होते हैं जिन्हें एक साथ बदलने की ज़रूरत होती है. ऐसा तब होता है, जब उन्हें अलग-अलग हिस्सों में बांटा गया हो. जैसे, पैड ऑपरेशन में पैडिंग डाइमेंशन.
• अन्य सभी फ़ैक्टर को पास-थ्रू फ़ैक्टर माना जाता है. इसका मतलब है कि जिन फ़ैक्टर को उनसे मैप किए गए सभी टेंसर में एक ही तरह से शर्ड किया जाता है उन्हें किसी भी तरह के कम्यूनिकेशन की ज़रूरत नहीं होती.

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

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

सीमाएं:

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

पैरामीटर:

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.

सीमाएं:

• 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
>

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

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

सीमाएं:

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

पैरामीटर:

TensorShardingPerValueAttr

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

सिंटैक्स:

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

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

सीमाएं:

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

पैरामीटर:

Enums

PropagationDirection

प्रोपगेशन डायरेक्शन एनम

उदाहरण: