StableHLO में डायनेमिज़म

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

डाइनैमिकिटी की शब्दावली और सहायता के बारे में खास जानकारी

सबसे पहले, इस दस्तावेज़ में दिखने वाले कुछ शब्दों के बारे में बताते हैं. साथ ही, StableHLO में इनकी सहायता के बारे में भी कम शब्दों में बताते हैं:

डाइनैमिक डाइमेंशन

डाइनैमिक डाइमेंशन का मतलब ऐसे डाइमेंशन से है जिसका डाइमेंशन साइज़ पता नहीं है. StableHLO में, हम ? यानी tensor<16x?xf32> का इस्तेमाल करके डाइनैमिक डाइमेंशन दिखाते हैं.

सीमित डाइनैमिक

सीमित डाइनैमिकिटी से ऐसे डाइनैमिक डाइमेंशन का पता चलता है जिसकी वैल्यू का ऊपरी सीमित वैल्यू पता हो. आम तौर पर, यह एक्ज़ीक्यूशन के दौरान टेंसर को पैडिंग के लिए काम का होता है. StableHLO में, हम #stablehlo.bounds का इस्तेमाल करके सीमित डाइनैमिकिटी को टेंसर कोडिंग के तौर पर दिखाते हैं. इसका मतलब है कि रैंक-2 टेंसर, जिसमें एक डाइनैमिक डाइमेंशन 16 तक सीमित है और दूसरे डाइमेंशन की कोई सीमा नहीं है, उसे tensor<?x?xf32, #stablehlo.bounds<16, ?>> के तौर पर दिखाया जा सकता है.

StableHLO, सीमित डाइनैमिकिटी दिखा सकता है. हालांकि, फ़्रेमवर्क के लिए सीमित सहायता उपलब्ध है. यह सहायता, TensorFlow से शुरू हुई और PyTorch/XLA में कुछ सहायता के साथ उपलब्ध है.

बेहिसाब डायनेमिज़म

नाम के मुताबिक, अनबाउंड डायनेमिज़्म का मतलब ऐसे डाइनैमिक डाइमेंशन से है जिसके साइज़ की कोई जानकारी नहीं है. JAX, PyTorch/XLA, और TF की सुविधा वाले StableHLO में, इस तरह की डायनेमिज़्म आम तौर पर होता है. इसे अक्सर डाइनैमिक बैच साइज़ या सीक्वेंस लंबाई वाले मॉडल एक्सपोर्ट करने के लिए इस्तेमाल किया जाता है.

StableHLO में, हम डाइनैमिकिटी के इस फ़ॉर्म के लिए, बाउंड कोडिंग को हटा देते हैं. जैसे, tensor<?x?xf32>.

शेप पॉलीमॉर्फ़िज़्म

शेप पॉलीमरफ़िज़्म एक ऐसा शब्द है जो हमें JAX से मिला है.

पॉलीमरफ़िज़्म को आकार देने के दो मुख्य नतीजे होते हैं:

1. प्रोग्राम में मौजूद सभी डाइनैमिकिटी, उसके इनपुट आर्ग्युमेंट पर निर्भर करती है.
2. डाइनैमिकिटी सिर्फ़ टेंसर आकार से जुड़ी होती है, यानी कि यह डेटा पर निर्भर नहीं होती.

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

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

डेटा पर निर्भर डाइनैमिक

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

डेटा पर निर्भर कई डाइनैमिक ऑपरेशन को सीमित डाइनैमिकिटी का इस्तेमाल करके मॉडल किया जा सकता है. इसमें टेंसर के साइज़ की ऊपरी सीमा तय की जाती है और आम तौर पर हार्डवेयर, टेंसर पैडिंग की मदद से इसे लागू करता है. फ़िलहाल, PyTorch/XLA और TensorFlow में डेटा पर निर्भर डाइनैमिकिटी के लिए कुछ सहायता मिलती है. हालांकि, JAX फ़िलहाल उन ऑपरेशन को ट्रैक नहीं करता है जिनसे डेटा पर निर्भर डाइनैमिकिटी पैदा होती है.

डाइनैमिक डाइमेंशन वाले प्रोग्राम एक्सपोर्ट करना

डाइनैमिक बैच साइज़ या सीक्वेंस की लंबाई वाले प्रोग्राम एक्सपोर्ट करने का तरीका जानने के लिए, StableHLO के ट्यूटोरियल देखें:

• JAX ट्यूटोरियल > डाइनैमिक बैच साइज़ के साथ एक्सपोर्ट करना
• PyTorch/XLA ट्यूटोरियल > डाइनैमिक बैच साइज़ के साथ एक्सपोर्ट करना

डाइनैमिक प्रोग्राम को बेहतर बनाने के लिए कंपाइलर पास

डाइनैमिक पास पाइपलाइन हटाएं

आकार को बेहतर बनाने के लिए कुछ काम के पास हैं. आसानी से इन सभी पास को पास पाइपलाइन createStablehloRemoveDynamismPipeline में इकट्ठा किया जा सकता है:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

डाइनैमिकिटी को बेहतर बनाने के लिए अलग-अलग पास

अलग-अलग, आकार को बेहतर बनाने के लिए ये पास काम के होते हैं:

• stablehlo-refine-arguments, इनपुट आर्ग्युमेंट को कंक्रीट टेंसर टाइप से बदलने के लिए.
• stablehlo-refine-shapes, ताकि पूरे प्रोग्राम में नए इनपुट आर्ग्युमेंट के आकार की जानकारी को प्रसारित किया जा सके.
• stablehlo-canonicalize-dynamism, डाइनैमिक ऑप्ट-आउट को उनके स्टैटिक वैरिएंट से बदलने के लिए.

अप-टू-डेट जानकारी और उदाहरणों के लिए, लिंक किए गए दस्तावेज़ देखें.

उदाहरण: डाइनैमिकिटी कैसे काम करती है और इसका इस्तेमाल कैसे किया जा सकता है?

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

स्टैटिक add_one मॉडल

इसे दिखाने के लिए, हम नीचे दिए गए आसान add_one मॉडल का इस्तेमाल करेंगे:

def add_one(x):
  return x + 1

tensor<4xf32> का इस्तेमाल करके ट्रैक करने पर, हमें StableHLO प्रोग्राम का यह डेटा मिलेगा:

// File: add_one.mlir
func.func @add_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<4xf32>
  return %0 : tensor<4xf32>
}

यह मॉडल, tensor<4xf32> आकार वाले इनपुट आर्ग्युमेंट के लिए सिर्फ़ काम करेगा. अगर हमने कभी भी अपने बैच का साइज़ या क्रम की लंबाई बदली है, तो हमें सोर्स कोड को फिर से ट्रैक करना होगा और उसे StableHLO पर फिर से कम करना होगा. इस बात की कोई गारंटी नहीं है कि हमारे पास अब भी सोर्स कोड का ऐक्सेस है!

डाइनैमिक add_one मॉडल

यहां से आकार पॉलिमॉर्फ़िक डायनेमिज़्म का काम शुरू होता है. इसके बजाय, JAX और PyTorch/XLA, डाइनैमिक तौर पर मान्य आईआर के साथ add_one मॉडल को उत्सर्जित कर सकते हैं. यह डाइनैमिक इनपुट शेप से मैच करने के लिए, कॉन्स्टेंट को इस तरह ब्रॉडकास्ट करेगा:

// File: add_one_dynamic.mlir
func.func public @main(%arg0: tensor<?xf32>) -> tensor<?xf32> {
  %cst = stablehlo.constant dense<1.0> : tensor<f32>
  %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?xf32>) -> tensor<i32>
  %1 = stablehlo.reshape %0 : (tensor<i32>) -> tensor<1xi32>
  %2 = stablehlo.dynamic_broadcast_in_dim %cst, %1, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %3 = stablehlo.add %arg0, %2 : tensor<?xf32>
  return %3 : tensor<?xf32>
}

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

डाइनैमिक मॉडल को बेहतर बनाना

उदाहरण के लिए, पास के क्रम को इस तरह सेट करके, इस प्रोग्राम को बेहतर बनाया जा सकता है:

stablehlo-opt add_one_dynamic.mlir \
  --stablehlo-refine-arguments='types=tensor<16xf32>' \
  --stablehlo-refine-shapes \
  --stablehlo-canonicalize-dynamism

धीरे-धीरे, इस तरह से प्रोग्राम में बदलाव होता है:

// After stablehlo-refine-arguments: Inputs updated, shapes not propagated
func.func public @main(%arg0: tensor<16xf32>) -> tensor<?xf32> {
  %c = stablehlo.constant dense<16> : tensor<1xi64>
  %0 = stablehlo.custom_call @stablehlo.shape_refinement_operand_wrapper(%arg0, %c) {indices_of_shape_operands = dense<1> : tensor<1xi64>} : (tensor<16xf32>, tensor<1xi64>) -> tensor<?xf32>
  ...
  %3 = stablehlo.dynamic_broadcast_in_dim %cst, %2, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<?xf32>
  %4 = stablehlo.add %0, %3 : tensor<?xf32>
  return %4 : tensor<?xf32>
}

// After stablehlo-refine-shapes: Shapes propagated, dynamic ops still exist
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %c = stablehlo.constant dense<16> : tensor<1xi32>
  %0 = stablehlo.dynamic_broadcast_in_dim %cst, %c, dims = [] : (tensor<f32>, tensor<1xi32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// After stablehlo-canonicalize-dynamism: Dynamic ops replaced with static ops
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
  %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
  %1 = stablehlo.add %arg0, %0 : tensor<16xf32>
  return %1 : tensor<16xf32>
}

// (Bonus) Use ` --stablehlo-aggressive-simplification` pass to canonicalize the
// constant broadcast, leaving us with the original static program in this case.
func.func public @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
  %cst = stablehlo.constant dense<1.000000e+00> : tensor<16xf32>
  %0 = stablehlo.add %arg0, %cst : tensor<16xf32>
  return %0 : tensor<16xf32>
}