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

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

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

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

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

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

डाइनैमिक कॉन्टेंट दिखाने की सुविधा

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

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

अनबाउंड डाइनैमिज़्म

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

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

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

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

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

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

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

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

डेटा के आधार पर डाइनैमिक तरीके से बदलाव करने की सुविधा

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

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

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

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

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

डाइनैमिज़्म पास पाइपलाइन हटाना

आकृतियों को बेहतर बनाने के लिए, कुछ काम के पास उपलब्ध हैं. ये सभी पास, पास पाइपलाइन createStablehloRemoveDynamismPipeline में बंडल किए गए हैं:

void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
                                           TypeRange refinedTypes);

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

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

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

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

उदाहरण: डाइनैमिज़्म कैसे फ़ायदेमंद है और मैं इसका इस्तेमाल कैसे करूं?

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

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