StableHLO 最初是從 MHLO 方言啟動,並且繼承類型推論的 MHLO 實作。您可以在 status.md 中追蹤實作進度。
以下建議準則旨在確保 StableHLO 運算實作高品質的驗證器和形狀函式。
建議採行的做法
這些提案適用於重新審視現有的實作方式,並達到新的營運目標,直到全面涵蓋為止。
(P1) 使用 StableHLO 規格做為可靠資料來源
spec是 StableHLO 運算的所有驗證器和形狀函式的可靠資料來源。每個運算結果的現有驗證器和形狀函式都需要重新審視,才能與規範完全一致。請注意,規格文件會不斷演進。如果無法取得運算的規格,請改用 XLA 實作做為可靠資料來源,包括 xla/service/shape_inference.cc 和 xla/service/hlo_verifier.cc。XLA 實作並未涵蓋不受限的動態生成,因此在無限制的動態生成技術在動態式 RFC 出現前,我們會以常識判斷。
(P2) 充分運用 ODS
ODS 檔案 (例如 StablehloOps.td) 會以特徵和類型為每個運算元/屬性/結果定義運算,並執行驗證。因此,對於 ODS 保證的屬性,驗證器或形狀函式並不需要驗證碼。如果與 ODS 重複,請移除驗證碼,因為系統絕不會觸發驗證碼。
我們需要針對 ODS 中的限制條件新增測試嗎?請參閱「建立測試規範」。
(P3) 維護驗證器和圖形函式中的驗證碼
兩者:
- 驗證器:由
Op::verify()
實作,以及 - 形狀函式:由
Op::inferReturnTypes()
或Op::inferReturnTypeComponents
等InferTypeOpInterface
實作
可能擁有驗證碼來檢查運算元/屬性/結果。初始分割作業可能如下所示:讓驗證器檢查運算元/屬性,然後讓形狀函式只計算推測結果類型,然後檢查實際結果類型的相容性。但實際上這種拆分方式有幾個問題:
- 自動產生的
build()
函式可以呼叫形狀函式,而不必先呼叫驗證器。因此,相關輸入內容也必須在形狀函式中驗證。 - 程式碼重複:例如,在驗證器中,我們會對運算元進行一些處理,然後驗證一些中繼結果。然後在形狀函式中,這些中繼結果有助於推斷最終結果。這些中繼結果必須計算兩次。
- 維護負擔:營運驗證包含在兩種不同方法中。
解決方法如下:
針對沒有地區的大部分運算 (例如
PadOp
):請嘗試將所有驗證碼放入形狀函式中,並完全捨棄驗證器。如果因無法推斷傳回類型 (例如使用ReshapeOp
或BroadcastInDimOp
) 而無法實現,請建立含有必要驗證邏輯的驗證器。一般而言,可推論的作業 (例如AddOp
) 可能還是需要驗證器來執行其他驗證,因為這些作業是驗證所提供傳回類型的限制,而這些作業無法在類型/形狀推論方法中存取。針對使用區域的運算 (例如
ReduceOp/IfOp
;完整清單請見這裡):自動產生的建構工具不會使用區域做為參數,因此如果這些建構工具涉及類型推論,系統就會以空白區域呼叫形狀函式 (請參閱此範例)。如果不需要地區推論類型推論 (例如
ReduceOp
),請將區域相關驗證邏輯放在驗證器中,而不是形狀函式。如果不可避免,請重複部分程式碼。如果類型推論需要區域 (
IfOp/CaseOp/MapOp
),那麼形狀函式必須明確驗證區域並非空白,即使 ODS 可能已保證在 Op 定義中存在該區域。
(P4) 建立測試規範
是否需要對 ODS 涵蓋的驗證進行新增/維護測試?
我們不會。測試應專注於驗證器和形狀函式,而對 ODS 所做的變更需要重新查看此操作。
但請謹慎處理遺漏的片段:舉例來說,如果運算包含特徵 SameOperandsAndResultShape
,該特徵只會檢查形狀,而不會檢查元素類型,則運算元類型/結果的元素類型驗證仍需測試。
我們在哪裡執行驗證和類型推論的測試?
ops_stablehlo.mlir 包含運算的陽性案例,以及 (至少) 每個驗證錯誤執行 1 個負測試。也能確認推測的傳回類型與實際結果類型是否「相容」 (並非與實際結果類型相同)。
infer_stablehlo.mlir 會透過 hlo_test_infer.get_return_type_components"(%x):...
行驗證運算的形狀函式是否存在,並檢查推論類型是否完全符合預期。一般來說,每項運算需執行一次陽性測試。
進行方式
實作或重新查看運算的驗證器和/或形狀函式時:
請將所有正面案例和負面情況放入 ops_stablehlo.mlir。
在 infer_stablehlo.mlir 中加入單一陽性測試以測試介面。
(選用) 如果運算作業很複雜,且可能包含大量測試,請考慮在同一個資料夾中新增名為
verify_<op_name>.mlir
或verify_<your_topic>.mlir
的獨立測試檔案。