型別推斷

StableHLO 最初是從 MHLO 方言啟動,並且繼承類型推論的 MHLO 實作。您可以在 status.md 中追蹤實作進度。

以下建議準則旨在確保 StableHLO 運算實作高品質的驗證器和形狀函式。

建議採行的做法

這些提案適用於重新審視現有的實作方式,並達到新的營運目標,直到全面涵蓋為止。

(P1) 使用 StableHLO 規格做為可靠資料來源

spec是 StableHLO 運算的所有驗證器和形狀函式的可靠資料來源。每個運算結果的現有驗證器和形狀函式都需要重新審視,才能與規範完全一致。請注意,規格文件會不斷演進。如果無法取得運算的規格,請改用 XLA 實作做為可靠資料來源,包括 xla/service/shape_inference.ccxla/service/hlo_verifier.cc。XLA 實作並未涵蓋不受限的動態生成,因此在無限制的動態生成技術在動態式 RFC 出現前,我們會以常識判斷。

(P2) 充分運用 ODS

ODS 檔案 (例如 StablehloOps.td) 會以特徵和類型為每個運算元/屬性/結果定義運算,並執行驗證。因此,對於 ODS 保證的屬性,驗證器或形狀函式並不需要驗證碼。如果與 ODS 重複,請移除驗證碼,因為系統絕不會觸發驗證碼。

我們需要針對 ODS 中的限制條件新增測試嗎?請參閱「建立測試規範」。

(P3) 維護驗證器和圖形函式中的驗證碼

兩者:

  • 驗證器:由 Op::verify() 實作,以及
  • 形狀函式:由 Op::inferReturnTypes()Op::inferReturnTypeComponentsInferTypeOpInterface 實作

可能擁有驗證碼來檢查運算元/屬性/結果。初始分割作業可能如下所示:讓驗證器檢查運算元/屬性,然後讓形狀函式只計算推測結果類型,然後檢查實際結果類型的相容性。但實際上這種拆分方式有幾個問題:

  • 自動產生的 build() 函式可以呼叫形狀函式,而不必先呼叫驗證器。因此,相關輸入內容也必須在形狀函式中驗證。
  • 程式碼重複:例如,在驗證器中,我們會對運算元進行一些處理,然後驗證一些中繼結果。然後在形狀函式中,這些中繼結果有助於推斷最終結果。這些中繼結果必須計算兩次。
  • 維護負擔:營運驗證包含在兩種不同方法中。

解決方法如下:

  1. 針對沒有地區的大部分運算 (例如 PadOp):請嘗試將所有驗證碼放入形狀函式中,並完全捨棄驗證器。如果因無法推斷傳回類型 (例如使用 ReshapeOpBroadcastInDimOp) 而無法實現,請建立含有必要驗證邏輯的驗證器。一般而言,可推論的作業 (例如 AddOp) 可能還是需要驗證器來執行其他驗證,因為這些作業是驗證所提供傳回類型的限制,而這些作業無法在類型/形狀推論方法中存取。

  2. 針對使用區域的運算 (例如 ReduceOp/IfOp;完整清單請見這裡):自動產生的建構工具不會使用區域做為參數,因此如果這些建構工具涉及類型推論,系統就會以空白區域呼叫形狀函式 (請參閱此範例)。

    1. 如果不需要地區推論類型推論 (例如 ReduceOp),請將區域相關驗證邏輯放在驗證器中,而不是形狀函式。如果不可避免,請重複部分程式碼。

    2. 如果類型推論需要區域 (IfOp/CaseOp/MapOp),那麼形狀函式必須明確驗證區域並非空白,即使 ODS 可能已保證在 Op 定義中存在該區域。

(P4) 建立測試規範

是否需要對 ODS 涵蓋的驗證進行新增/維護測試?

我們不會。測試應專注於驗證器和形狀函式,而對 ODS 所做的變更需要重新查看此操作。

但請謹慎處理遺漏的片段:舉例來說,如果運算包含特徵 SameOperandsAndResultShape,該特徵只會檢查形狀,而不會檢查元素類型,則運算元類型/結果的元素類型驗證仍需測試。

我們在哪裡執行驗證和類型推論的測試?

ops_stablehlo.mlir 包含運算的陽性案例,以及 (至少) 每個驗證錯誤執行 1 個負測試。也能確認推測的傳回類型與實際結果類型是否「相容」 (並非與實際結果類型相同)。

infer_stablehlo.mlir 會透過 hlo_test_infer.get_return_type_components"(%x):... 行驗證運算的形狀函式是否存在,並檢查推論類型是否完全符合預期。一般來說,每項運算需執行一次陽性測試。

進行方式

實作或重新查看運算的驗證器和/或形狀函式時:

  1. 請將所有正面案例和負面情況放入 ops_stablehlo.mlir

  2. infer_stablehlo.mlir 中加入單一陽性測試以測試介面。

  3. (選用) 如果運算作業很複雜,且可能包含大量測試,請考慮在同一個資料夾中新增名為 verify_<op_name>.mlirverify_<your_topic>.mlir 的獨立測試檔案。