अनुवादक डिज़ाइन

डेटा मॉडल

StableHLO प्रोग्राम, टेंसर (एन-डाइमेंशन ऐरे) पर आधारित कंप्यूटेशन होते हैं. इन्हें मौजूदा मॉडल में Tensor क्लास का इस्तेमाल करके लागू किया जाता है. Tensor ऑब्जेक्ट के लिए मौजूदा स्टोरेज क्लास, detail::Buffer, टेंसर के mlir::ShapedType और mlir::HeapAsmResourceBlob ऑब्जेक्ट को एकmlir::HeapAsmResourceBlob ऑब्जेक्ट के साथ सेव करती है. यह टेंसर डेटा के म्यूटेबल ब्लॉब के तौर पर, मुख्य क्रम में तय किए गए बाइट ऐरे के तौर पर दिखती है. मेमोरी मैनेजमेंट को आसान बनाने के लिए, detail::Buffer ऑब्जेक्ट को रेफ़रंस के तौर पर गिना जाता है.

टेंसर के अलग-अलग एलिमेंट को Element क्लास का इस्तेमाल करके दिखाया जाता है. इसमें स्टोरेज के लिए, APInt, APFloat या pair<APFloat,APFloat> में से एक को पकड़े हुए विभेदित यूनियन का इस्तेमाल किया जाता है. आखिरी वाले का इस्तेमाल मुश्किल टाइप वाले एलिमेंट को स्टोर करने के लिए किया जाता है.

Tensor के अलग-अलग एलिमेंट से इंटरैक्ट करने के लिए, यहां दिए गए एपीआई मौजूद हैं:

• Element Tensor::get(llvm::ArrayRef<int64_t> index): मल्टी-डाइमेंशनल इंडेक्स index में Element ऑब्जेक्ट के तौर पर, किसी एक टेंसर एलिमेंट को एक्सट्रैक्ट करने के लिए.
• void Tensor::set(llvm::ArrayRef<int64_t> index, Element element);: मल्टी-डाइमेंशन इंडेक्स index पर, Element ऑब्जेक्ट element को टेंसर में अपडेट करने के लिए.

अनुवादक कैसे काम करता है

अनुवादक के लिए एंट्री फ़ंक्शन यह है

SmallVector<Tensor> eval(func::FuncOp func, ArrayRef<Tensor> args);

जिससे ये काम करते हैं:

1. यह func के SSA आर्ग्युमेंट और उससे जुड़े रनटाइम Tensor की वैल्यू को ट्रैक करता है. ये वैल्यू, args में दिए गए हैं. इसके लिए, सिंबल टेबल मैप, M का इस्तेमाल किया जाता है.
2. func के हर ऑपरेटर के लिए, SSACFG के क्रम में:
   • ऑपरेटर की मदद से, eval को शुरू करता है. ऑपरेटर के हर SSA ऑपरेंड के लिए, M से उसका रनटाइम वैल्यू एक्सट्रैक्ट करें, ताकि उसे eval को शुरू करने के लिए तर्क के तौर पर इस्तेमाल किया जा सके.
   • ऑप के SSA नतीजे और M में आकलन की गई वैल्यू ट्रैक करता है.

(2) में बताए गए ओपी-लेवल eval, ऑपरेटर की एक्ज़ीक्यूशन सिमैंटिक लागू करने की ज़िम्मेदारी है. stablehlo::AddOp का एक उदाहरण नीचे दिया गया है. इस उदाहरण में, lhs और rhs टेंसर के अलग-अलग एलिमेंट को, जोड़ी के हिसाब से Element ऑब्जेक्ट के तौर पर निकाला गया है. इसके बाद, एलिमेंट जोड़े गए. जोड़ने के नतीजे, Element ऑब्जेक्ट को आखिरी result टेंसर में सेव किया जाता है.

Tensor eval(AddOp op, const Tensor &lhs, const Tensor &rhs) {
  Tensor result(op.getType());

  for (auto it = result.index_begin(); it != result.index_end(); ++it)
    result.set(*it, lhs.get(*it) + rhs.get(*it));

  return result;
}

कुल मिलाकर देखें, तो अनुवादक के डिज़ाइन को इस तरह ऑप्टिमाइज़ किया गया है कि अलग-अलग ऑपरेशन के लिए, eval फ़ंक्शन लागू किए जा सकें. ऐसा इसलिए है, क्योंकि इसका मकसद StableHLO के लिए एक रेफ़रंस लागू करना है. उदाहरण के लिए, eval को टेंप्लेट फ़ंक्शन के तौर पर तय करने और एलिमेंट टाइप के साथ पैरामीटर सेट करने के बजाय, हम Element::operator+ वगैरह में अलग-अलग तरह के एलिमेंट को कैसे हैंडल करते हैं, इसकी जानकारी शामिल करते हैं. इससे, eval को आसानी से लागू किया जा सकता है.

कॉन्सटेंट फ़ोल्डिंग के लिए इंटरप्रेटर का इस्तेमाल करना

हम कॉन्सटैंट ऑपरेंड वैल्यू के साथ ऑपरेशन को फ़ोल्ड करने के लिए, इंटरप्रेटर सिस्टम का इस्तेमाल कर सकते हैं. यह कोड स्निपेट, फ़्लोटिंग-पॉइंट टाइप किए गए ऑपरेंड के साथ stablehlo::AddOp को फ़ोल्ड करने के आइडिया को दिखाता है:

OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
  auto attrs = adaptor.getOperands();
  DenseElementsAttr lhsData = attrs[0].dyn_cast<DenseElementsAttr>();
  DenseElementsAttr rhsData = attrs[1].dyn_cast<DenseElementsAttr>();
  if (!lhsData || !rhsData) return {};

  auto lhs = Tensor(lhsData);
  auto rhs = Tensor(rhsData);
  auto result = eval(*this, lhs, rhs);

  SmallVector<APFloat> values;
  for (auto i = 0; i < result.getNumElements(); ++i) {
    Element element = result.get(i);
    values.push_back(element.getValue().cast<FloatAttr>().getValue());
  }

  return DenseElementsAttr::get(result.getType(), values);
}

फ़िलहाल, हम अनुवादक को लगातार फ़ोल्डिंग में जोड़ने पर काम नहीं कर रहे हैं, क्योंकि हम StableHLO के लिए फ़ोल्डर लागू नहीं करना चाहते हैं. हालांकि, आने वाले समय में, हम MHLO में कॉन्सटैंट फ़ोल्डिंग के लिए इंटरप्रेटर का इस्तेमाल करने की योजना बना रहे हैं.ऐसा करने पर, हम ऊपर दिए गए कोड स्निपेट के अर्गोनॉमिक्स को बेहतर बनाएंगे. उदाहरण के लिए, हमारे पास ऐसा हेल्पर फ़ंक्शन हो सकता है जो Tensor ऑब्जेक्ट में कॉन्सटैंट ऑपरेंड को पैक करता है और Tensor नतीजों को OpFoldResult में अनपैक करता है.

StableHLO अनुवादक की जांच करना

इंटरप्रेटर, इनपुट के तौर पर (A) एक StableHLO प्रोग्राम और (B) डेटा की वैल्यू, प्रोग्राम को फ़ीड करने के लिए लेता है. साथ ही, आउटपुट डेटा वैल्यू जनरेट करता है, जिन्हें उपयोगकर्ता से मिले डेटा की वैल्यू से मैच किया जाता है. डेटा वैल्यू (B) को प्रोग्राम में ही, stablehlo.constant ऑपरेशन का इस्तेमाल करके हार्ड कोड किया जाता है. अनुवादक, इनपुट प्रोग्राम की जांच करता है. टेस्ट के तहत आने वाले ऑप के आउटपुट की जांच जांच के ज़रिए की जाती है (उदाहरण के लिए, check.expect_eq, check.expect_almost_eq), जैसा कि नीचे दिखाया गया है. check.expect_eq और check.expect_eq_const, यह देखें कि इस्तेमाल किए जा सकने वाले किसी भी टाइप के लिए, बिट के आधार पर बराबरी की क्वालिटी है या नहीं. check.expect_almost_eq और check.expect_almost_eq_const फ़्लोटिंग पॉइंट और कॉम्प्लेक्स टाइप के लिए, टेस्टिंग से जुड़े दिशा-निर्देशों (G6) में दी गई सहिष्णुता में, बराबरी की समानता की जांच करें.

// CHECK-LABEL: Evaluated results of function: add_op_test_ui4
func.func @add_op_test_ui4() {
  %0 = stablehlo.constant dense<[0, 2]> : tensor<2xui4>
  %1 = stablehlo.constant dense<[15, 3]> : tensor<2xui4>
  %2 = stablehlo.add %0, %1 : tensor<2xui4>
  check.expect_eq_const %2, [15, 5] : tensor<2xui4>
  func.return
}

टेस्ट यूटिलिटी stablehlo-translate --interpret (कोड) प्रोग्राम को पार्स करने के लिए ज़िम्मेदार है. यह फ़ंक्शन, हर फ़ंक्शन के साथ-साथ फ़ंक्शन बनाने वाले ऑपरेशन की जानकारी भी देता है. हमारे पास एक खास टेस्ट-सुइट है, जिसमें हर StableHLO Op के लिए, कई रनटाइम व्यवहार वाले कई टेस्ट शामिल हैं. ये टेस्ट यहां देखे जा सकते हैं (जैसे कि interpret_*.mlir).

जांच के दिशा-निर्देश

(G1) क्या हमें हर ऑपरेटर के लिए, इस्तेमाल किए जा सकने वाले सभी टाइप की जांच करनी चाहिए?

हम निम्न नियमों के एक संयोजन का उपयोग करके यह तय कर सकते हैं:

1. किसी ऑपरेटर को लागू करते समय, अगर किसी खास टाइप को मैनेज करने के लिए, उससे जुड़े eval फ़ंक्शन में कोड मौजूद होता है, तो उस टाइप को लागू करने के लिए टेस्ट होना ज़रूरी है. उदाहरण के तौर पर, add ऑप के लिए पूर्णांक, बूलियन, फ़्लोटिंग-पॉइंट, और कॉम्प्लेक्स टाइप को मैनेज करने के लिए एक खास कोड होता है. इसलिए, हमें टाइप की हर कैटगरी के लिए एक टेस्ट की ज़रूरत होती है.

2. अगर टाइप के किसी सेट को, उससे जुड़े eval फ़ंक्शन में एक जैसे तरीके से हैंडल किया जाता है, तो उन सभी टाइप के लिए एक ही टेस्ट काफ़ी होना चाहिए. उदाहरण के तौर पर, addऑप के लिए, पूर्णांक टाइप (si4, u4, si8, u8 वगैरह) के सभी वैरिएंट को एक तरह से मैनेज किया जाता है. इसके लिए llvm::APInt एपीआई का इस्तेमाल किया जाता है. इसलिए, हम हर वैरिएंट के लिए टेस्ट जोड़ने की प्रोसेस को छोड़कर, एक रिप्रज़ेंटेटिव टेस्ट जोड़ सकते हैं. प्रतिनिधि को चुनते समय कोई गड़बड़ी न हो, इसके लिए हमें इन दिशा-निर्देशों का पालन करना चाहिए:

   • अगर सभी टाइप को एक जैसा रखा जाता है और उनका प्रिमिटिव टाइप एक ही है (जैसे, अगर सभी पूर्णांक या फ़्लोटिंग-पॉइंट या कॉम्प्लेक्स टाइप हैं), तो सबसे ज़्यादा बिट-विथ वाली वैल्यू चुनें.
   • अगर सभी टाइप को एक समान तरीके से मैनेज किया जाता है, तो इनमें प्रिमिटिव टाइप शामिल हैं. इसके बाद, प्राथमिकता के घटते क्रम में, नीचे दिए गए प्रिमिटिव टाइप वाले विकल्प को चुनें: पूर्णांक, फ़्लोटिंग-पॉइंट, बूलियन, कॉम्प्लेक्स.

(G2) हम यह कैसे तय करेंगे कि किसी ऑपरेटर के व्यवहार को कवर करने के लिए, कुल कितने टेस्ट ज़रूरी हैं?

इसका मकसद, जांच की कम से कम संख्या के साथ, Op के लिए इंटरप्रेटर के लॉजिक को बेहतर तरीके से शामिल करना (यानी लागू करने के सभी कोने के केस) शामिल करना है. रखरखाव के लिए, टेस्ट की संख्या कम करना ज़रूरी है. हमारे पास जितने कम टेस्ट हैं उनकी समीक्षा करना उतना ही आसान है. साथ ही, यह पक्का करना भी आसान है कि वे ऑपरेटर को विस्तार से समझाते हैं. इस वजह से, हमें उम्मीद है कि आसान ऑपरेशन के लिए सिर्फ़ एक ही टेस्ट उपलब्ध होगा. अगर किसी वजह से, अगर किसी वजह से कवरेज सही तरीके से नहीं हो पा रहा है, तो उसे 90% से ज़्यादा पर भी रोका जा सकता है. पुल के अनुरोध की समीक्षा करते समय, हर मामले के हिसाब से इस पर फ़ैसला लिया जाएगा.

(G3) अनुवादक इन्फ़्रास्ट्रक्चर के लिए टेस्ट जोड़ने के बारे में क्या ख्याल है?

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

(G4) अगर किसी ऑपर्च्यूनिटी को लागू करना दूसरे ऑपरेशन पर निर्भर करता है, तो क्या हमें बाद वाले ऑपरेशन के लिए टेस्ट लिखने चाहिए?

नहीं. उदाहरण के लिए, batch_norm_grad को लागू करने की प्रोसेस divide, subtract, multiply, और अन्य चीज़ों के हिसाब से हो सकती है. हमें पुराने ऑपरेशन की जांच करके, जांच करने से बचना चाहिए.

(G5) क्या हमें तय किए गए / तय न किए गए व्यवहार को लागू करने के लिए, टेस्ट लिखने चाहिए?

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

(G6) फ़्लोटिंग-पॉइंट टाइप के लिए टेस्ट लिखते समय, जांच में अनुमानित नतीजे को किस सटीक जानकारी देने की ज़रूरत होती है?

प्राइमरी ऑपरेशन (जोड़ना, घटाना, गुणा करना, भाग देना, और स्क्वेयर करना) के लिए, आईईईई स्पेसिफ़िकेशन को लागू करने से यह उम्मीद की जाती है कि वह, गणित के हिसाब से सटीक नतीजे के 0.5 यूएलपी के दायरे में पूरा नतीजा मिले. इसलिए, हम सुरक्षित तरीके से कल्पना कर सकते हैं कि इन कार्रवाइयों से मिलने वाले नतीजे, ज़्यादा से ज़्यादा एक यूएलपी के बराबर होंगे. हालांकि, ऐसा हो सकता है कि यह उन ट्रांसेंडेंटल फ़ंक्शन (sine, cosine वगैरह) के लिए काम न करे जिनके लिए, सटीक गारंटी वाले फ़ंक्शन के बारे में बताया गया है (वजह).

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

func.func @check_tolerance() {
  %0 = stablehlo.constant dense<0.2> : tensor<f32>

  // The following check succeeds as %0 is almost equal to the provided
  // constant modulo the tolerance, mentioned above.
  check.expect_almost_eq_const %0, dense<0.19999> : tensor<f32>

  // The following check fails as %0 is not bitwise equal to the provided
  // constant.
  check.expect_eq_const %0, dense<0.19999> : tensor<f32>

  func.return
}

यह StableHLO ऑपरेशन की संख्या के सटीक होने की जांच करने का सिर्फ़ पहला कदम है. फ़िलहाल, StableHLO की खास जानकारी में कम शब्दों में जानकारी दी गई है. #1156 का पता लगाने के लिए, StableHLO का इस्तेमाल करने के हमारे अनुभव और हिस्सेदारों से मिले सुझाव के आधार पर काम जारी है. जैसे-जैसे यह प्रक्रिया आगे बढ़ेगी, हम इन्फ़्रास्ट्रक्चर को इसके हिसाब से अपडेट करेंगे.

(G7) टेस्ट की कोडिंग स्टाइल के बारे में कुछ जानना है?

1. पक्का करें कि SSA वैल्यू (जैसे कि %0, %1 वगैरह) में डिफ़ॉल्ट तौर पर सेट करने के बजाय, इनपुट/आउटपुट का असली नाम इस्तेमाल किया जाए
2. अगर टेस्ट मौजूद हैं, तो पक्का करें कि प्रिटी-प्रिंट वाले फ़ॉर्मैट का इस्तेमाल किया गया हो.

(G8) क्या हमें स्पेसिफ़िकेशन में, पहले से दिए गए उदाहरण को शामिल करना चाहिए? हां (टेस्टिंग की पूरी जानकारी के लिए).