自主資料型別：在TVM中啟用自定義資料型別探索
介紹
在設計加速器時，一個重要的決定是如何在硬體中近似地表示實數。這個問題有一個長期的行業標準解決方案：IEEE 754浮點標準.1。然而，當試圖通過構建高度專業化的設計來最大限度地利用硬體時，使用通用IEEE 754浮點有意義嗎？如果知道工作負載的數字需求，是否可以構建一個更小、更快或更省電的資料型別？答案是肯定的！研究人員已經開始在學術和工業加速器設計中嘗試新的資料型別。例如，Google的張量處理單元（TPU）使用bfloat型別：一個被截斷為16位的單精度IEEE浮點。由於許多深度學習工作負載對數值要求不嚴格，這種截斷通常對模型精度沒有影響，同時立即將儲存成本降低一半。
然而，在研究人員開始為資料型別構建硬體之前，首先需要確定資料型別在工作負載中的數值行為。這通常需要首先構建其資料型別的軟體模擬版本（例如Berkeley SoftFloat或libposit），然後將資料型別直接插入工作負載，以檢視工作負載如何使用資料型別執行。更好的方法是將資料型別直接整合到編譯器本身中，這樣就可以編譯許多不同的工作負載來使用該資料型別。這兩種路徑都可能是乏味的，鑑於現代編譯器的規模和複雜性，後一種路徑往往變得不可管理。一個來自GitHub的例子顯示有人將posit資料型別侵入TensorFlow。結果是237次提交，新增了近6000行程式碼，並在程式碼庫中處理了200多個檔案，這僅僅是新增了一個資料型別！對許多研究人員來說，這樣的工作量是令人望而卻步的。
為了解決這些問題，提出了自帶資料型別框架。該框架允許使用者將模擬的資料型別插入TVM，從而可以輕鬆地在深度學習工作負載中探索新的資料型別。與上面的Tensorflow示例中的positions不同，它支援編譯器中的單個新資料型別，而Bring Your Own datatype框架支援各種各樣的使用者定義型別。
自主資料型別
“自主資料型別”框架的目標是使使用者能夠使用自定義資料型別執行深入學習的工作負載。在Bring Your Own Datatypes框架中，“datatype”表示標量型別：例如float或uint。不處理像Intel Flexpoint這樣複雜的資料塊格式（floating point or Intel’s Flexpoint）。此外，只宣告支援這些標量資料型別的軟體模擬版本；不明確支援在自定義資料型別硬體上編譯和執行。
TVM中的每個張量都分配了一個型別程式碼，它定義了張量中標量的資料型別。這些型別程式碼在TVM中具有硬編碼的含義，對映到諸如int和float之類的常見資料型別。然而，絕大多數型別程式碼都是未使用的。“自帶資料型別”框架允許使用者宣告這些未使用的型別程式碼，並在執行時新增自己的新資料型別。
該框架是作為登錄檔來實現的，它位於TVM的常規資料型別架構上。使用者與資料型別登錄檔互動的主要方式有兩種：第一種是資料型別註冊，第二種是降低函式註冊。這些步驟分別類似於資料型別的宣告和實現。
本文中引用的所有程式碼都基於TVM儲存庫的主分支提交4cad71d。將使用一個示例posit資料型別，該資料型別位於src/target/datatype/posit/posit-wrapper.cc並且可以在TVM中使用USE_BYODT_POSIT標誌進行編譯。
Datatype Registration
為了註冊資料型別，使用者為資料型別分配一個名稱和一個型別程式碼，其中型別程式碼來自自定義資料型別可用的未使用型別程式碼範圍。
tvm.target.datatype.register(‘posit’, 150)
The above code registers the ‘posit’ datatype with type code 150. This registration step allows TVM to parse programs which use the custom type:
x = relay.var(‘x’, shape=(3, ), dtype=‘float32’)
y = relay.var(‘y’, shape=(3, ), dtype=‘float32’)
x_posit = relay.cast(x, dtype=‘custom[posit]16’)
y_posit = relay.cast(y, dtype=‘custom[posit]16’)
z_posit = x_posit + y_posit
z = relay.cast(z_posit, dtype=‘float32’)
program = relay.Function([x, y], z)
print(program)

v0.0.4

fn (%x: Tensor[(3), float32], %y: Tensor[(3), float32]) {

%0 = cast(%x, dtype=“custom[posit]16”);

%1 = cast(%y, dtype=“custom[posit]16”);

%2 = add(%0, %1);

cast(%2, dtype=“float32”)

}

上面的程式將float32輸入x和y轉換成posits，相加，然後將結果轉換回float32。一旦posit型別被註冊，TVM就能夠解析特殊的dtype語法custom[]，其中是為型別註冊的名稱。此語法還支援通常的x格式；這裡，使用16來表示每個posite的寬度是16位。（通道數預設為1。）
Lowering Function Registration
雖然TVM可以解析上述程式，但它還不能編譯它，因為TVM還不知道如何在posit型別上編譯操作。為了編譯這些程式，為自定義資料型別註冊了降低函式，這有助於TVM將操作轉換成它可以理解和編譯的內容。
一般來說，使用者不需要直接將操作降低到LLVM或CUDA。相反，大多數使用自定義資料型別的程式碼可以通過一些簡單的技巧被簡化為不使用自定義資料型別的程式碼。然後，可以依賴於本機TVM來理解和編譯程式碼。

圖1：使用者註冊的降低函式的預期結果。
降低函式應該將使用自定義資料型別的程式轉換為本機TVM可以理解和編譯的程式（在本例中，呼叫外部庫，使用兩個uint16）。
圖1顯示了一個通用模式。假設對探索posit型別感興趣，並選擇通過將posit模擬庫（例如Stillwater Universal）通過自帶的資料型別框架插入TVM來執行一些工作負載。工作量是一個簡單的程式，它新增了兩個posite輸入。Native-TVM不知道如何實現posit加法，但它不需要，因為有一個實現資料型別的庫！這個庫包含一個正數加法的實現，以及乘法和平方根等其他運算子。為了實現這個posit加法，只想呼叫庫。因此，Add節點應該成為一個呼叫節點，呼叫庫中的一個函式（稱之為position16es2add）。為了將輸入位置的位儲存在TVM理解的型別中，使用16位無符號整數。生成的程式是TVM可以理解和編譯的程式，它只是呼叫一個外部庫函式，取兩個無符號整數。
為了實現上述降低，為posit註冊了一個降低功能：
tvm.target.datatype.register_op(
tvm.target.datatype.create_lower_func({16: ‘Posit16es2Add’}),
‘Add’, ‘llvm’, ‘posit’)
上面的程式碼為特定運算子（Add）、編譯目標（LLVM）、資料型別（posit）和位長度（16）註冊了一個降低函式。第一個引數是lowering函式。這可以是任何獲取TVM IR節點並返回新TVM IR節點的函式。在例子中，使用了Bring Your Own Datatypes框架提供的helper函式。tvm.target.datatype.create_lower_func({16:‘Posit16es2Add’})為上述公共模式建立一個降低函式。結果函式將給定節點的引數轉換為uint16_t，然後將節點本身轉換為對給定函式名的呼叫（在本例中，對於位長度為16的position16es2add）。我們傳遞一個字典來create_lower_func，以便TVM可以根據資料型別的位長度將其分派到適當的函式名。
要實現自定義資料型別，使用者將需要為他們希望執行的工作負載中的每個運算子註冊一個降低函式。對於像ResNet這樣的網路，大約有10個操作符，包括Add、Div、各種型別轉換和Max。在測試中，註冊一個資料型別和所有的降低函式需要大約40行Python。一旦註冊了所有需要的運算元，自定義資料型別工作負載就可以像任何其他TVM程式一樣輕鬆執行！
結束
“自帶資料型別”框架將使用者定義的資料型別引入TVM。希望這將鼓勵資料型別研究人員在研究中使用TVM；同樣，希望這將激發深度學習社群中對定製資料型別的興趣。有關“自帶資料型別”框架的更多文件，請訪問“將自己的資料型別帶到TVM開發人員”。