在鵬城實驗室實習的時候,主要的工作是開發AI算子標準參考實現。參考實現是對算子標準的補充,旨在提供代碼級別的語義的規范,因此正確性是它的最重要屬性。
開發一個新的算子接口是比較繁瑣的,雖然標準文檔中對算子接口已有比較明確的定義,但是為了給予標準實現者最大的優化空間,它沒有規定實現的細節,這也給參考實現的開發增加了一點難度,因為我在開發新算子的時候常常要對這個算子重新調研,找到它的具體實現邏輯,再開發對應的接口。為了滿足正確性的要求,在開發完一個新算子后,需要對它進行單元測試。這時候應盡可能給出不同情況的樣例,然后將參考實現的結果與tensorflow、pytorch等框架的結果進行對比。
開發了一些算子接口以后,我發現不同算子之間會存在大量重復代碼,經常有復雜算子包含了簡單算子代碼的情況出現,比如maxpool會重復max的代碼。為了減少代碼開發的工作量、降低代碼維護難度,我先開發簡單算子,然后讓復雜算子中重復了簡單算子的代碼的部分替換成直接調用簡單算子的方式。這種開發方式可能對性能要求高的算子庫不適用,但是參考實現追求的是正確性,不考慮性能,因此參考實現的所有算子都可以采取這種方法進行開發。后來,參考實現的維護難度也降低了不少。
除了參考實現,我們還開發了一套自動化測試框架。對于不同的算子庫來說,只要遵循了算子標準,那么它們對應的算子接口應該是差不多的:返回值是狀態碼,輸入輸出都在參數列表中,參數之間的前后順序也一致。對算子進行正確性測試的邏輯也是一致的,即生成輸入樣例、算子運算、對比結果。對于計算機來說,不同的地方在于數據結構(張量等)和函數名。那么,我們只要把測試樣例和測試流程代碼提前寫好,讓用戶把自行開發的算子的函數名和創建張量等數據結構的接口注冊到自動化測試框架中,框架就可以自動地對算子進行正確性測試了。所以,自動化測試框架的關鍵技術就是接口注冊機制。我們調研以后發現,Google Test的“類型參數化”機制可以讓用戶在觸發測試的時候才指定測試代碼中某些變量的數據類型。我們就把“類型參數化”作為接口注冊機制的原型,通過宏把用戶的接口封裝到一個類中,用戶的接口就“假扮”成了一種數據類型,然后我們在測試代碼中從對應的變量里提取出用戶的接口,就可以實現接口注冊了。
完成以上工作以后,我發現最大的工作量不是寫代碼,而是寫代碼之前的調研和設計。我也曾試過草草調研之后就著手寫代碼,短期內似乎是可行的,但是在后期只要遇到一個bug,我就得花大量時間去調試、調研,得不償失。運籌帷幄之中,決勝千里之外。
作者:黎子毅