指南：從注冊到Lowering全流程)
1. 為什么要在 Relay 里加算子從一次真實需求說起深度學習編譯器這個領域Relay 作為 TVM 的高層中間表示承擔著把前端框架的計算圖翻譯成可優(yōu)化、可調(diào)度 IR 的核心職責。日常做模型部署的時候絕大多數(shù)算子 Relay 的算子庫都已經(jīng)覆蓋了卷積、矩陣乘、各類激活函數(shù)、歸一化、池化這些常規(guī)操作基本都能直接找到對應實現(xiàn)。但真正落到項目里總會遇到一些“標準庫里沒有”的情況——可能是某個論文里新提出的激活函數(shù)可能是業(yè)務側自定義的融合邏輯也可能是某個冷門框架導出的算子 Relay 沒有對應解析。這時候繞不過去的一件事就是自己給 Relay 添加一個新算子。這篇文章面向的讀者是已經(jīng)跑通過 TVM 基礎流程、能看懂 Relay IR、知道relay.op大概長什么樣但還沒完整走過“注冊一個新算子”全流程的人。我會把從算子定義、屬性注冊、類型關系推導、到 lowering 落地、再到 Python 側調(diào)用和測試的整條鏈路拆開講清楚。整個過程涉及 C 和 Python 兩側的配合涉及 TVM 的注冊機制、類型系統(tǒng)、以及 Relay 的 pass 體系信息量不小但走通一遍之后你會發(fā)現(xiàn)它的套路其實非常固定。先說清楚一個前提給 Relay 加算子本質(zhì)上是在做三件事。第一告訴 Relay 這個算子長什么樣——它的名字、輸入輸出、有哪些屬性參數(shù)第二告訴 Relay 這個算子的類型規(guī)則——給定輸入類型輸出類型怎么推第三告訴 Relay 這個算子最終怎么算——也就是 lowering 到更低層的實現(xiàn)。這三件事分別對應算子注冊、類型關系TypeRelation、以及計算實現(xiàn)。很多人卡住不是因為某一步特別難而是因為這三件事分散在不同文件、不同注冊宏里第一次做的時候容易顧此失彼。我個人的建議是動手之前先想清楚你的算子屬于哪一類。如果它可以用現(xiàn)有 Relay 算子組合表達出來那優(yōu)先考慮用relay.op里的組合方式實現(xiàn)而不是新增算子——組合方式不需要碰 C維護成本低得多。只有當組合方式表達不了、或者性能上必須下沉到 C 甚至要自定義 schedule 的時候才值得走完整的新增算子流程。這個判斷很關鍵我見過不少人一上來就加算子結果發(fā)現(xiàn)用te.compute加幾個現(xiàn)有 op 就能搞定白白增加了維護負擔。下面我按實際開發(fā)順序把整個流程拆成幾個部分來講。每一部分我都會說明“為什么這么做”而不只是“怎么做”因為 TVM 的注冊機制設計得比較繞不理解意圖的話照抄代碼很容易在細節(jié)上翻車。2. 動手前的整體設計新算子該放在哪一層2.1 先判斷組合實現(xiàn)還是新增算子在決定新增算子之前有一個必須做的判斷這個算子能不能用 Relay 已有的算子拼出來。舉個常見的例子假設你要實現(xiàn)一個swish激活函數(shù)它的定義是x * sigmoid(x)。這個完全可以用relay.multiply(x, relay.sigmoid(x))表達不需要新增任何算子。再比如某些自定義的歸一化如果能拆成 reduce、subtract、divide 的組合也優(yōu)先走組合路線。組合路線的優(yōu)勢非常明顯不碰 C不需要重新編譯 TVMPython 側直接就能用而且能自動享受 Relay 現(xiàn)有的融合、常量折疊等優(yōu)化。缺點是如果組合出來的子圖比較碎可能會影響融合效果或者在某些后端上性能不理想。什么時候必須新增算子我總結了幾種典型情況。第一種是算子語義無法用現(xiàn)有算子組合表達比如涉及特殊的內(nèi)存布局或者不可分解的數(shù)學運算。第二種是性能敏感組合實現(xiàn)會產(chǎn)生大量中間張量而新增算子可以在一個 compute 里完成。第三種是需要自定義 schedule比如你想對這個算子做特殊的循環(huán)切分和向量化組合實現(xiàn)沒法控制底層調(diào)度。第四種是前端框架導出的模型里帶了這個算子你需要讓 Relay 能解析它。判斷清楚之后如果確定要新增接下來要決定算子的“層級”。Relay 里算子可以注冊在幾個不同的層次最上層是 Relay oprelay.op命名空間它直接面向計算圖往下是 TOPITVM Operator Inventory提供具體的 compute 和 schedule 實現(xiàn)再往下是 TETensor Expression描述張量計算。一個完整的算子通常需要在 Relay 層注冊 op在 TOPI 層提供 compute 實現(xiàn)必要時還要寫 schedule。2.2 文件組織新增算子該改哪些文件TVM 的代碼組織比較清晰新增一個算子通常涉及以下幾個位置。C 側算子注冊一般放在src/relay/op/下的對應子目錄里比如張量運算放tensor/神經(jīng)網(wǎng)絡相關放nn/。類型關系通常和算子注冊放在一起或者單獨放在op/op_common.h之類的公共頭文件里。TOPI 的 compute 實現(xiàn)放在python/tvm/topi/下對應的模塊里比如topi/nn/、topi/math/。Python 側的封裝放在python/tvm/relay/op/下。這里有個經(jīng)驗不要一上來就想著把算子加到最合適的位置先把流程跑通再考慮代碼放哪。我見過有人花大量時間糾結目錄結構結果算子還沒跑起來。正確的做法是先在一個臨時文件里把注冊、類型關系、compute 都寫出來編譯通過、測試通過之后再按 TVM 的代碼規(guī)范挪到合適的位置。另外要提醒的是TVM 的注冊宏是全局的算子名字在整個 Relay 命名空間里必須唯一。命名的時候建議遵循現(xiàn)有算子的命名習慣比如nn.contrib_xxx、math.xxx這種帶命名空間前綴的形式避免和已有算子沖突。如果你只是本地實驗可以用一個不太可能沖突的前綴比如myop.xxx。2.3 注冊機制的核心三個宏各管什么TVM 里新增算子最核心的是三個注冊宏理解它們各自的作用是走通流程的關鍵。第一個是RELAY_REGISTER_OP它注冊算子的“元信息”算子名字、參數(shù)描述、類型關系、以及一些屬性。這個宏決定了 Relay 在解析和類型推導階段怎么看待這個算子。它里面會指定set_attrs_type、add_argument、set_support_level、set_attr等。第二個是TVM_REGISTER_GLOBAL它注冊一個全局函數(shù)通常用來把算子的 compute 實現(xiàn)暴露給 Relay 的 lowering 階段。Relay 在把 op 轉成 TE 的時候會通過這個全局函數(shù)名去調(diào)用對應的 compute。第三個是類型關系相關的注冊比如RelayExpr TypeRel或者用TypeRelation的子類。它負責在類型推導階段計算輸出類型。對于大多數(shù)逐元素算子可以直接用現(xiàn)成的BroadcastRel或者ElemwiseRel對于形狀有特殊變化的算子就需要自己寫一個 TypeRelation。這三個宏的關系可以這樣理解RELAY_REGISTER_OP是“身份證”告訴系統(tǒng)有這么個算子類型關系是“規(guī)則書”告訴系統(tǒng)輸入輸出類型怎么對應TVM_REGISTER_GLOBAL是“施工隊”告訴系統(tǒng)實際怎么算。三者缺一不可而且名字要對應上否則會在不同階段報錯。3. 核心細節(jié)拆解算子注冊與類型關系3.1 用 RELAY_REGISTER_OP 描述算子元信息先看算子注冊。一個典型的注冊代碼長這樣以我實際做過的一個自定義激活算子為例這里用虛構的myop.leaky_clamp代稱RELAY_REGISTER_OP(myop.leaky_clamp) .describe(Rcode(Leaky clamp activation. Computes max(min(x, upper), lower) with a leaky slope for the lower part. )code TVM_ADD_FILELINE) .set_num_inputs(1) .add_argument(data, Tensor, The input tensor.) .set_support_level(10) .set_attrs_typeLeakyClampAttrs() .add_type_rel(LeakyClamp, LeakyClampRel());這里逐項說明。describe是給文檔和調(diào)試用的描述寫清楚算子語義方便別人包括未來的自己理解。set_num_inputs(1)聲明輸入個數(shù)注意這里只算數(shù)據(jù)輸入屬性參數(shù)不算在內(nèi)。add_argument描述每個輸入的名字和類型這個名字會出現(xiàn)在錯誤信息和調(diào)試輸出里。set_support_level是 TVM 內(nèi)部的一個標記數(shù)值越大表示支持程度越低自定義算子一般給個較大的值比如 10表示“非官方支持”。set_attrs_type指定屬性結構體類型如果你的算子有額外參數(shù)比如上面的lower、upper、slope就需要定義一個 Attrs 結構體。add_type_rel綁定類型關系第一個參數(shù)是類型關系的名字第二個是實例。屬性結構體通常這樣定義struct LeakyClampAttrs : public tvm::AttrsNodeLeakyClampAttrs { double lower; double upper; double slope; TVM_DECLARE_ATTRS(LeakyClampAttrs, relay.attrs.LeakyClampAttrs) { TVM_ATTR_FIELD(lower).set_default(-1.0).describe(Lower bound.); TVM_ATTR_FIELD(upper).set_default(1.0).describe(Upper bound.); TVM_ATTR_FIELD(slope).set_default(0.01).describe(Leaky slope.); } };這個結構體繼承AttrsNode用TVM_DECLARE_ATTRS聲明字段每個字段用TVM_ATTR_FIELD注冊并可以設默認值。屬性結構體的好處是 Relay 會自動幫你做參數(shù)解析、類型檢查、以及序列化。定義好之后Python 側調(diào)用時傳的關鍵字參數(shù)會自動映射到這些字段。這里有個容易踩的坑屬性字段的類型要選對。TVM 支持int、double、bool、String、Array等類型但如果你用了不支持的復雜類型編譯會報錯。另外屬性結構體的名字上面那個字符串relay.attrs.LeakyClampAttrs要全局唯一建議帶上算子名前綴。3.2 類型關系輸出類型到底怎么推類型關系是很多人第一次做時最容易迷糊的地方。它的職責是給定輸入表達式的類型推導出輸出表達式的類型。對于 Relay 來說類型主要是TensorType包含 shape 和 dtype。最簡單的做法是復用現(xiàn)成的類型關系。如果你的算子是逐元素的輸出 shape 和輸入完全一致直接用ElemwiseRel就行。如果涉及廣播比如兩個不同 shape 的輸入做運算用BroadcastRel。這兩個現(xiàn)成的類型關系覆蓋了相當一部分場景。但有些算子形狀會變比如 reduce 類算子輸出維度會減少或者像 reshape 這種輸出 shape 由屬性決定。這時候就需要自己寫 TypeRelation。一個自定義 TypeRelation 的骨架class LeakyClampRel : public TypeRelation { public: LeakyClampRel() : TypeRelation(nullptr, LeakyClampRel) {} TVM_DLL static ArrayType LeakyClampRelFunc( const ArrayType types, int num_inputs, const Attrs attrs, const TypeReporter reporter) { ICHECK_EQ(types.size(), 2); const auto* data types[0].asTensorTypeNode(); if (data nullptr) return types; reporter-Assign(types[1], TensorType(data-shape,>tvm.script.register_object def leaky_clamp(data, lower-1.0, upper1.0, slope0.01): return _make.leaky_clamp(data, lower, upper, slope)這里的_make.leaky_clamp是 TVM 根據(jù) C 注冊自動生成的參數(shù)順序要和 C 側add_argument加屬性字段的順序?qū)?。實際上 TVM 的 FFI 機制會自動把 Python 調(diào)用轉成 C 的Call節(jié)點屬性會打包成 Attrs。更規(guī)范的做法是用tvm.relay.op的注冊裝飾器把算子注冊到relay.op命名空間下這樣用戶可以用relay.op.myop.leaky_clamp(...)調(diào)用。同時建議加上類型檢查比如檢查輸入是不是 Tensor屬性值是否在合理范圍內(nèi)。Python 側還有一個重要工作是注冊算子的“梯度”或者“策略”如果你的算子需要參與自動微分或者有特定的算子融合策略。梯度注冊用register_gradient融合策略用register_fusion_pattern。這些不是必須的但如果你希望新算子能融入 Relay 的優(yōu)化流程最好補上。4. 實操過程從 compute 到 lowering 的完整鏈路4.1 用 TOPI 寫 compute 實現(xiàn)算子注冊和類型關系搞定之后接下來是實際的計算實現(xiàn)。Relay 的 lowering 會把 op 轉成 TE而 TE 的 compute 通常由 TOPI 提供。TOPI 的 compute 函數(shù)接收輸入張量和屬性返回一個 TE 的Tensor。繼續(xù)用leaky_clamp舉例TOPI 實現(xiàn)大概是這樣def leaky_clamp(data, lower-1.0, upper1.0, slope0.01): def _compute(ins): x ins[0] return te.compute( x.shape, lambda *idx: te.max( te.min(x[idx], tvm.tir.const(upper, x.dtype)), tvm.tir.const(lower, x.dtype) ) tvm.tir.const(slope, x.dtype) * te.min( x[idx] - tvm.tir.const(lower, x.dtype), 0.0 ), nameleaky_clamp, tagelemwise ) return _compute這里有幾個細節(jié)值得說。第一te.compute的第一個參數(shù)是輸出 shape逐元素算子直接用輸入 shape。第二lambda 里的idx是索引變量x[idx]取出對應位置的元素。第三常量要用tvm.tir.const包裝并指定 dtype否則可能因為類型不匹配報錯。第四tag參數(shù)很重要它標記了算子的類型比如elemwise、broadcast、reduceRelay 的融合 pass 會根據(jù) tag 決定能不能把這個算子和相鄰算子融合。第五name參數(shù)會出現(xiàn)在生成的代碼里取個有意義的名字方便調(diào)試。寫完 compute 之后需要用TVM_REGISTER_GLOBAL把它注冊成全局函數(shù)函數(shù)名要和 Relay 側 lowering 時查找的名字一致。通常的約定是topi.myop.leaky_clamp這種形式。Relay 在 lowering 時會根據(jù)算子的名字拼出這個全局函數(shù)名去調(diào)用。這里有個非常容易踩的坑compute 函數(shù)的參數(shù)順序。它必須和 Relay 側add_argument加屬性字段的順序完全一致否則運行時會因為參數(shù)錯位產(chǎn)生莫名其妙的錯誤。我建議在寫的時候把參數(shù)列表寫下來對照不要憑記憶。4.2 注冊 compute 與 lowering 對接compute 寫好后注冊代碼通常長這樣TVM_REGISTER_GLOBAL(topi.myop.leaky_clamp) .set_body([](TVMArgs args, TVMRetValue* rv) { *rv LeakyClampCompute(args[0], args[1], args[2], args[3]); });或者更常見的是在 Python 側用tvm.register_func注冊tvm.register_func(topi.myop.leaky_clamp) def _leaky_clamp_compute(data, lower, upper, slope): return leaky_clamp(data, lower, upper, slope)兩種方式都可以C 側注冊性能稍好Python 側注冊更靈活、改起來不用重新編譯。開發(fā)階段我強烈建議先用 Python 注冊因為改一行代碼就能生效不用等 C 編譯。等邏輯穩(wěn)定了再考慮挪到 C。lowering 對接的關鍵是 Relay 怎么知道去調(diào)用哪個全局函數(shù)。Relay 的OpImplementation或者FTVMCompute屬性會指定 compute 函數(shù)名。在RELAY_REGISTER_OP里可以用set_attrFTVMCompute(FTVMCompute, ...)來指定或者依賴默認的命名約定。默認約定是算子名加前綴比如myop.leaky_clamp對應topi.myop.leaky_clamp。如果你用了非默認命名一定要顯式設置FTVMCompute。4.3 編譯、測試與驗證代碼寫完之后編譯 TVM。如果你只改了 Python 側不需要重新編譯直接跑測試就行。如果改了 C 側需要重新編譯這個過程可能比較久建議用ninja增量編譯。測試分幾個層次。第一層是算子能不能被正確解析和類型推導寫個簡單的 Relay 程序用relay.Function構造調(diào)用跑relay.transform.InferType看輸出類型對不對。第二層是數(shù)值正確性構造隨機輸入用新算子算一遍再用 NumPy 手動實現(xiàn)算一遍對比結果。第三層是端到端把算子放進一個完整模型編譯到目標后端跑推理看結果。數(shù)值測試的代碼大概這樣import numpy as np import tvm from tvm import relay data relay.var(data, relay.TensorType((2, 3), float32)) out relay.op.myop.leaky_clamp(data, lower-1.0, upper1.0, slope0.01) func relay.Function([data], out) mod tvm.IRModule.from_expr(func) mod relay.transform.InferType()(mod) x np.random.randn(2, 3).astype(float32) expected np.maximum(np.minimum(x, 1.0), -1.0) 0.01 * np.minimum(x 1.0, 0.0) with tvm.transform.PassContext(opt_level3): lib relay.build(mod, targetllvm) runtime tvm.contrib.graph_executor.GraphModule(lib[default](tvm.cpu())) runtime.set_input(data, x) runtime.run() result runtime.get_output(0).numpy() np.testing.assert_allclose(result, expected, rtol1e-5)這個測試覆蓋了類型推導、lowering、編譯、執(zhí)行全流程。如果這一步過了說明算子基本可用。如果沒過錯誤信息通常會指向某個具體階段按階段排查即可。5. 常見問題與排查技巧實錄5.1 編譯期報錯注冊相關的典型問題新增算子時編譯期報錯主要集中在注冊環(huán)節(jié)。最常見的是“算子未注冊”或者“找不到類型關系”。前者通常是RELAY_REGISTER_OP的名字和調(diào)用時用的名字不一致注意 Relay 里算子名是大小寫敏感的而且命名空間前綴要完全匹配。后者是add_type_rel綁定的名字和實際 TypeRelation 注冊的名字不一致。另一個常見問題是屬性結構體沒注冊。如果你定義了LeakyClampAttrs但忘了用TVM_REGISTER_NODE_TYPE注冊編譯時可能不報錯但運行時解析屬性會失敗。注冊代碼是TVM_REGISTER_NODE_TYPE(LeakyClampAttrs);這個宏要放在屬性結構體定義之后。我踩過一次坑屬性結構體定義在頭文件里注冊宏忘了加結果算子能注冊成功但一調(diào)用就崩排查了半天才發(fā)現(xiàn)是屬性沒注冊。還有一種情況是set_num_inputs和實際add_argument的次數(shù)不匹配。set_num_inputs聲明的是數(shù)據(jù)輸入個數(shù)add_argument也要調(diào)用相同次數(shù)。如果聲明了 2 個輸入但只add_argument了 1 次類型推導時types數(shù)組長度會對不上報數(shù)組越界的錯。5.2 運行期報錯類型推導與 lowering 的坑運行期報錯里類型推導相關的最多。典型癥狀是Check failed: reporter-Assign之類的錯誤通常是因為輸出類型和輸入類型不兼容或者Assign被調(diào)用了多次。檢查你的 TypeRelation確保每個輸出只Assign一次且 dtype 和 shape 都合理。lowering 階段的報錯通常是“找不到 compute 函數(shù)”。這多半是全局函數(shù)名和 Relay 期望的名字不一致。排查方法是打印出 Relay 期望的函數(shù)名或者直接在RELAY_REGISTER_OP里顯式設置FTVMCompute屬性。另外compute 函數(shù)返回的必須是 TE 的 Tensor 或者 Tensor 列表如果返回了別的東西lowering 會失敗。數(shù)值不對但能跑通的情況通常是 compute 實現(xiàn)里的索引或者常量寫錯了。比如te.min和te.max的參數(shù)順序搞反或者常量 dtype 和輸入 dtype 不一致導致隱式類型轉換。這類問題建議先用小 shape 的輸入手動驗證比如(1, 1)或者(2, 2)把中間結果打印出來對比。5.3 常見問題速查表問題現(xiàn)象可能原因排查方向算子未注冊名字不匹配或注冊宏未執(zhí)行檢查RELAY_REGISTER_OP名字與調(diào)用名類型關系找不到add_type_rel名字與注冊名不一致核對 TypeRelation 注冊字符串屬性解析失敗屬性結構體未注冊補TVM_REGISTER_NODE_TYPE數(shù)組越界set_num_inputs與add_argument次數(shù)不符核對輸入個數(shù)聲明找不到 compute全局函數(shù)名不匹配顯式設置FTVMCompute或核對命名數(shù)值錯誤compute 索引或常量寫錯小 shape 手動驗證中間結果融合失效tag 設置不當檢查 compute 的 tag 參數(shù)5.4 獨家避坑經(jīng)驗第一條經(jīng)驗開發(fā)階段全部用 Python 注冊 compute不要碰 C。C 編譯一次動輒十幾分鐘改一個符號就要重編效率極低。等 Python 版本驗證通過再考慮要不要挪到 C。實際上很多自定義算子用 Python 注冊完全夠用性能差異在大多數(shù)場景下可以忽略。第二條經(jīng)驗類型關系先用現(xiàn)成的實在不行再自己寫。ElemwiseRel和BroadcastRel能覆蓋大部分逐元素和廣播場景自己寫 TypeRelation 容易在邊界情況上出錯。如果非寫不可先寫一個最簡版本只處理TensorType把nullptr和動態(tài) shape 的情況留到后面再補。第三條經(jīng)驗測試要從最小用例開始。不要一上來就用真實模型的復雜輸入先用(1, 1)或者(2, 3)這種小 shape把類型推導、lowering、數(shù)值都對一遍。小用例跑通了再逐步放大。我見過有人直接上大模型報錯信息淹沒在幾百行日志里根本定位不到問題。第四條經(jīng)驗善用relay.transform.InferType和relay.transform.PrintIR。類型推導出問題時先跑InferType看報錯位置lowering 出問題時用PrintIR把 lowering 前后的 IR 打出來對比能快速定位是哪一步出的問題。第五條經(jīng)驗屬性字段的默認值要合理。如果某個屬性沒有默認值用戶調(diào)用時必須顯式傳否則解析會失敗。給一個合理的默認值能降低使用門檻也能避免因為漏傳參數(shù)導致的詭異錯誤。6. 算子擴展后的優(yōu)化與融合考量6.1 讓新算子參與算子融合新算子加進去只是第一步能不能融入 Relay 的優(yōu)化流程才是決定性能的關鍵。Relay 的融合 pass 會把相鄰的算子合并成一個 kernel減少內(nèi)存訪問和 kernel 啟動開銷。新算子能不能被融合取決于它的 tag 和融合模式注冊。compute 里的tag參數(shù)是融合 pass 的重要依據(jù)。elemwise標簽的算子通常可以和前后算子自由融合broadcast次之reduce和injective有各自的融合規(guī)則。如果你希望新算子能融合tag 要設對。比如逐元素算子設elemwise涉及廣播的設broadcast。除了 tag還可以用register_fusion_pattern注冊更精細的融合模式。比如你希望某個算子和它前面的卷積融合可以寫一個模式匹配規(guī)則。這部分相對進階不是所有算子都需要但如果你的算子性能敏感值得花時間研究。6.2 自定義 schedule 的接入點如果新算子的 compute 用默認 schedule 性能不理想可以自定義 schedule。TOPI 的 compute 函數(shù)通常配套一個 schedule 函數(shù)用te.create_schedule創(chuàng)建調(diào)度然后對循環(huán)做切分、向量化、并行化等操作。自定義 schedule 的接入方式是在 Relay 的 lowering 配置里指定。TVM 的relay.build可以傳schedule_cfg或者用OpStrategy機制把 compute 和 schedule 綁定。對于自定義算子最簡單的做法是在 TOPI 里提供一個 schedule 函數(shù)然后在 lowering 時通過FTVMSchedule屬性指定。schedule 的編寫是另一個大話題這里不展開。但要提醒一點不要過早優(yōu)化 schedule。先用默認 schedule 把功能跑通測出性能瓶頸再針對性地寫 schedule。很多時候默認 schedule 已經(jīng)夠用自定義 schedule 的收益可能還不如換個更合適的 tag 讓融合 pass 發(fā)揮作用。6.3 跨后端兼容性檢查新算子寫完之后如果你需要它在多個后端上都能跑要做兼容性檢查。不同后端對算子的支持程度不一樣比如某些后端不支持特定的數(shù)據(jù)類型或者不支持動態(tài) shape。測試的時候至少覆蓋 CPU 和你有需求的目標后端。如果某個后端不支持通常有兩種處理方式。一種是提供 fallback 實現(xiàn)讓不支持的部分回退到 CPU 或者用其他算子組合。另一種是在 lowering 階段做后端特定的轉換把新算子拆解成后端支持的算子。具體選哪種取決于你的性能要求和實現(xiàn)成本。還有一個容易被忽略的點是 dtype 支持。你的 compute 實現(xiàn)里如果用了tvm.tir.const包裝常量要確保常量的 dtype 和輸入 dtype 一致。如果輸入可能是float16也可能是float32compute 要能處理兩種 dtype不能寫死。測試的時候把常見 dtype 都跑一遍避免上線后才發(fā)現(xiàn)某個 dtype 不支持。7. 我個人在實際操作中的幾點體會走完整個流程我最深的體會是TVM 的算子注冊機制雖然看起來繁瑣但它的設計是有內(nèi)在邏輯的。注冊、類型關系、compute 這三層分別對應“是什么”“怎么推”“怎么算”把這三層想清楚代碼就是水到渠成的事。真正難的不是寫代碼而是理解每一層為什么存在、什么時候需要自定義、什么時候可以復用。另一個體會是文檔和社區(qū)示例很重要但不能全信。TVM 版本迭代快不同版本的 API 可能有差異網(wǎng)上找到的示例代碼不一定適用于你的版本。最可靠的做法是直接看當前版本 TVM 源碼里已有的算子實現(xiàn)照葫蘆畫瓢。比如你想加一個逐元素算子就找一個現(xiàn)有的逐元素算子把它的注冊、類型關系、compute 都讀一遍然后照著改。這比看任何教程都管用。最后一點不要怕報錯。新增算子涉及的環(huán)節(jié)多第一次做幾乎不可能一次通過。報錯信息通常很具體指向某個文件某一行順著排查基本都能找到原因。我自己的習慣是每改一小步就編譯測試一次不要攢一大堆改動再一起測那樣出了問題很難定位。小步快跑逐步驗證是走通這個流程最穩(wěn)的方式。如果你后續(xù)還想擴展可以考慮的方向包括給算子加梯度支持讓它能參與訓練、寫自定義 schedule 提升性能、或者把算子貢獻回 TVM 社區(qū)。每個方向都有各自的坑但核心流程和這篇文章講的是一致的。把基礎流程走通后面的擴展就是在這個框架上做加法。 TVM Relay自定義算子開發(fā)指南：從注冊到Lowering全流程
1. 為什么要在 Relay 里加算子從一次真實需求說起深度學習編譯器這個領域Relay 作為 TVM 的高層中間表示承擔著把前端框架的計算圖翻譯成可優(yōu)化、可調(diào)度 IR 的核心職責。日常做模型部署的時候絕大多數(shù)算子 Relay 的算子庫都已經(jīng)覆蓋了卷積、矩陣乘、各類激活函數(shù)、歸一化、池化這些常規(guī)操作基本都能直接找到對應實現(xiàn)。但真正落到項目里總會遇到一些“標準庫里沒有”的情況——可能是某個論文里新提出的激活函數(shù)可能是業(yè)務側自定義的融合邏輯也可能是某個冷門框架導出的算子 Relay 沒有對應解析。這時候繞不過去的一件事就是自己給 Relay 添加一個新算子。這篇文章面向的讀者是已經(jīng)跑通過 TVM 基礎流程、能看懂 Relay IR、知道relay.op大概長什么樣但還沒完整走過“注冊一個新算子”全流程的人。我會把從算子定義、屬性注冊、類型關系推導、到 lowering 落地、再到 Python 側調(diào)用和測試的整條鏈路拆開講清楚。整個過程涉及 C 和 Python 兩側的配合涉及 TVM 的注冊機制、類型系統(tǒng)、以及 Relay 的 pass 體系信息量不小但走通一遍之后你會發(fā)現(xiàn)它的套路其實非常固定。先說清楚一個前提給 Relay 加算子本質(zhì)上是在做三件事。第一告訴 Relay 這個算子長什么樣——它的名字、輸入輸出、有哪些屬性參數(shù)第二告訴 Relay 這個算子的類型規(guī)則——給定輸入類型輸出類型怎么推第三告訴 Relay 這個算子最終怎么算——也就是 lowering 到更低層的實現(xiàn)。這三件事分別對應算子注冊、類型關系TypeRelation、以及計算實現(xiàn)。很多人卡住不是因為某一步特別難而是因為這三件事分散在不同文件、不同注冊宏里第一次做的時候容易顧此失彼。我個人的建議是動手之前先想清楚你的算子屬于哪一類。如果它可以用現(xiàn)有 Relay 算子組合表達出來那優(yōu)先考慮用relay.op里的組合方式實現(xiàn)而不是新增算子——組合方式不需要碰 C維護成本低得多。只有當組合方式表達不了、或者性能上必須下沉到 C 甚至要自定義 schedule 的時候才值得走完整的新增算子流程。這個判斷很關鍵我見過不少人一上來就加算子結果發(fā)現(xiàn)用te.compute加幾個現(xiàn)有 op 就能搞定白白增加了維護負擔。下面我按實際開發(fā)順序把整個流程拆成幾個部分來講。每一部分我都會說明“為什么這么做”而不只是“怎么做”因為 TVM 的注冊機制設計得比較繞不理解意圖的話照抄代碼很容易在細節(jié)上翻車。2. 動手前的整體設計新算子該放在哪一層2.1 先判斷組合實現(xiàn)還是新增算子在決定新增算子之前有一個必須做的判斷這個算子能不能用 Relay 已有的算子拼出來。舉個常見的例子假設你要實現(xiàn)一個swish激活函數(shù)它的定義是x * sigmoid(x)。這個完全可以用relay.multiply(x, relay.sigmoid(x))表達不需要新增任何算子。再比如某些自定義的歸一化如果能拆成 reduce、subtract、divide 的組合也優(yōu)先走組合路線。組合路線的優(yōu)勢非常明顯不碰 C不需要重新編譯 TVMPython 側直接就能用而且能自動享受 Relay 現(xiàn)有的融合、常量折疊等優(yōu)化。缺點是如果組合出來的子圖比較碎可能會影響融合效果或者在某些后端上性能不理想。什么時候必須新增算子我總結了幾種典型情況。第一種是算子語義無法用現(xiàn)有算子組合表達比如涉及特殊的內(nèi)存布局或者不可分解的數(shù)學運算。第二種是性能敏感組合實現(xiàn)會產(chǎn)生大量中間張量而新增算子可以在一個 compute 里完成。第三種是需要自定義 schedule比如你想對這個算子做特殊的循環(huán)切分和向量化組合實現(xiàn)沒法控制底層調(diào)度。第四種是前端框架導出的模型里帶了這個算子你需要讓 Relay 能解析它。判斷清楚之后如果確定要新增接下來要決定算子的“層級”。Relay 里算子可以注冊在幾個不同的層次最上層是 Relay oprelay.op命名空間它直接面向計算圖往下是 TOPITVM Operator Inventory提供具體的 compute 和 schedule 實現(xiàn)再往下是 TETensor Expression描述張量計算。一個完整的算子通常需要在 Relay 層注冊 op在 TOPI 層提供 compute 實現(xiàn)必要時還要寫 schedule。2.2 文件組織新增算子該改哪些文件TVM 的代碼組織比較清晰新增一個算子通常涉及以下幾個位置。C 側算子注冊一般放在src/relay/op/下的對應子目錄里比如張量運算放tensor/神經(jīng)網(wǎng)絡相關放nn/。類型關系通常和算子注冊放在一起或者單獨放在op/op_common.h之類的公共頭文件里。TOPI 的 compute 實現(xiàn)放在python/tvm/topi/下對應的模塊里比如topi/nn/、topi/math/。Python 側的封裝放在python/tvm/relay/op/下。這里有個經(jīng)驗不要一上來就想著把算子加到最合適的位置先把流程跑通再考慮代碼放哪。我見過有人花大量時間糾結目錄結構結果算子還沒跑起來。正確的做法是先在一個臨時文件里把注冊、類型關系、compute 都寫出來編譯通過、測試通過之后再按 TVM 的代碼規(guī)范挪到合適的位置。另外要提醒的是TVM 的注冊宏是全局的算子名字在整個 Relay 命名空間里必須唯一。命名的時候建議遵循現(xiàn)有算子的命名習慣比如nn.contrib_xxx、math.xxx這種帶命名空間前綴的形式避免和已有算子沖突。如果你只是本地實驗可以用一個不太可能沖突的前綴比如myop.xxx。2.3 注冊機制的核心三個宏各管什么TVM 里新增算子最核心的是三個注冊宏理解它們各自的作用是走通流程的關鍵。第一個是RELAY_REGISTER_OP它注冊算子的“元信息”算子名字、參數(shù)描述、類型關系、以及一些屬性。這個宏決定了 Relay 在解析和類型推導階段怎么看待這個算子。它里面會指定set_attrs_type、add_argument、set_support_level、set_attr等。第二個是TVM_REGISTER_GLOBAL它注冊一個全局函數(shù)通常用來把算子的 compute 實現(xiàn)暴露給 Relay 的 lowering 階段。Relay 在把 op 轉成 TE 的時候會通過這個全局函數(shù)名去調(diào)用對應的 compute。第三個是類型關系相關的注冊比如RelayExpr TypeRel或者用TypeRelation的子類。它負責在類型推導階段計算輸出類型。對于大多數(shù)逐元素算子可以直接用現(xiàn)成的BroadcastRel或者ElemwiseRel對于形狀有特殊變化的算子就需要自己寫一個 TypeRelation。這三個宏的關系可以這樣理解RELAY_REGISTER_OP是“身份證”告訴系統(tǒng)有這么個算子類型關系是“規(guī)則書”告訴系統(tǒng)輸入輸出類型怎么對應TVM_REGISTER_GLOBAL是“施工隊”告訴系統(tǒng)實際怎么算。三者缺一不可而且名字要對應上否則會在不同階段報錯。3. 核心細節(jié)拆解算子注冊與類型關系3.1 用 RELAY_REGISTER_OP 描述算子元信息先看算子注冊。一個典型的注冊代碼長這樣以我實際做過的一個自定義激活算子為例這里用虛構的myop.leaky_clamp代稱RELAY_REGISTER_OP(myop.leaky_clamp) .describe(Rcode(Leaky clamp activation. Computes max(min(x, upper), lower) with a leaky slope for the lower part. )code TVM_ADD_FILELINE) .set_num_inputs(1) .add_argument(data, Tensor, The input tensor.) .set_support_level(10) .set_attrs_typeLeakyClampAttrs() .add_type_rel(LeakyClamp, LeakyClampRel());這里逐項說明。describe是給文檔和調(diào)試用的描述寫清楚算子語義方便別人包括未來的自己理解。set_num_inputs(1)聲明輸入個數(shù)注意這里只算數(shù)據(jù)輸入屬性參數(shù)不算在內(nèi)。add_argument描述每個輸入的名字和類型這個名字會出現(xiàn)在錯誤信息和調(diào)試輸出里。set_support_level是 TVM 內(nèi)部的一個標記數(shù)值越大表示支持程度越低自定義算子一般給個較大的值比如 10表示“非官方支持”。set_attrs_type指定屬性結構體類型如果你的算子有額外參數(shù)比如上面的lower、upper、slope就需要定義一個 Attrs 結構體。add_type_rel綁定類型關系第一個參數(shù)是類型關系的名字第二個是實例。屬性結構體通常這樣定義struct LeakyClampAttrs : public tvm::AttrsNodeLeakyClampAttrs { double lower; double upper; double slope; TVM_DECLARE_ATTRS(LeakyClampAttrs, relay.attrs.LeakyClampAttrs) { TVM_ATTR_FIELD(lower).set_default(-1.0).describe(Lower bound.); TVM_ATTR_FIELD(upper).set_default(1.0).describe(Upper bound.); TVM_ATTR_FIELD(slope).set_default(0.01).describe(Leaky slope.); } };這個結構體繼承AttrsNode用TVM_DECLARE_ATTRS聲明字段每個字段用TVM_ATTR_FIELD注冊并可以設默認值。屬性結構體的好處是 Relay 會自動幫你做參數(shù)解析、類型檢查、以及序列化。定義好之后Python 側調(diào)用時傳的關鍵字參數(shù)會自動映射到這些字段。這里有個容易踩的坑屬性字段的類型要選對。TVM 支持int、double、bool、String、Array等類型但如果你用了不支持的復雜類型編譯會報錯。另外屬性結構體的名字上面那個字符串relay.attrs.LeakyClampAttrs要全局唯一建議帶上算子名前綴。3.2 類型關系輸出類型到底怎么推類型關系是很多人第一次做時最容易迷糊的地方。它的職責是給定輸入表達式的類型推導出輸出表達式的類型。對于 Relay 來說類型主要是TensorType包含 shape 和 dtype。最簡單的做法是復用現(xiàn)成的類型關系。如果你的算子是逐元素的輸出 shape 和輸入完全一致直接用ElemwiseRel就行。如果涉及廣播比如兩個不同 shape 的輸入做運算用BroadcastRel。這兩個現(xiàn)成的類型關系覆蓋了相當一部分場景。但有些算子形狀會變比如 reduce 類算子輸出維度會減少或者像 reshape 這種輸出 shape 由屬性決定。這時候就需要自己寫 TypeRelation。一個自定義 TypeRelation 的骨架class LeakyClampRel : public TypeRelation { public: LeakyClampRel() : TypeRelation(nullptr, LeakyClampRel) {} TVM_DLL static ArrayType LeakyClampRelFunc( const ArrayType types, int num_inputs, const Attrs attrs, const TypeReporter reporter) { ICHECK_EQ(types.size(), 2); const auto* data types[0].asTensorTypeNode(); if (data nullptr) return types; reporter-Assign(types[1], TensorType(data-shape,>tvm.script.register_object def leaky_clamp(data, lower-1.0, upper1.0, slope0.01): return _make.leaky_clamp(data, lower, upper, slope)這里的_make.leaky_clamp是 TVM 根據(jù) C 注冊自動生成的參數(shù)順序要和 C 側add_argument加屬性字段的順序?qū)?。實際上 TVM 的 FFI 機制會自動把 Python 調(diào)用轉成 C 的Call節(jié)點屬性會打包成 Attrs。更規(guī)范的做法是用tvm.relay.op的注冊裝飾器把算子注冊到relay.op命名空間下這樣用戶可以用relay.op.myop.leaky_clamp(...)調(diào)用。同時建議加上類型檢查比如檢查輸入是不是 Tensor屬性值是否在合理范圍內(nèi)。Python 側還有一個重要工作是注冊算子的“梯度”或者“策略”如果你的算子需要參與自動微分或者有特定的算子融合策略。梯度注冊用register_gradient融合策略用register_fusion_pattern。這些不是必須的但如果你希望新算子能融入 Relay 的優(yōu)化流程最好補上。4. 實操過程從 compute 到 lowering 的完整鏈路4.1 用 TOPI 寫 compute 實現(xiàn)算子注冊和類型關系搞定之后接下來是實際的計算實現(xiàn)。Relay 的 lowering 會把 op 轉成 TE而 TE 的 compute 通常由 TOPI 提供。TOPI 的 compute 函數(shù)接收輸入張量和屬性返回一個 TE 的Tensor。繼續(xù)用leaky_clamp舉例TOPI 實現(xiàn)大概是這樣def leaky_clamp(data, lower-1.0, upper1.0, slope0.01): def _compute(ins): x ins[0] return te.compute( x.shape, lambda *idx: te.max( te.min(x[idx], tvm.tir.const(upper, x.dtype)), tvm.tir.const(lower, x.dtype) ) tvm.tir.const(slope, x.dtype) * te.min( x[idx] - tvm.tir.const(lower, x.dtype), 0.0 ), nameleaky_clamp, tagelemwise ) return _compute這里有幾個細節(jié)值得說。第一te.compute的第一個參數(shù)是輸出 shape逐元素算子直接用輸入 shape。第二lambda 里的idx是索引變量x[idx]取出對應位置的元素。第三常量要用tvm.tir.const包裝并指定 dtype否則可能因為類型不匹配報錯。第四tag參數(shù)很重要它標記了算子的類型比如elemwise、broadcast、reduceRelay 的融合 pass 會根據(jù) tag 決定能不能把這個算子和相鄰算子融合。第五name參數(shù)會出現(xiàn)在生成的代碼里取個有意義的名字方便調(diào)試。寫完 compute 之后需要用TVM_REGISTER_GLOBAL把它注冊成全局函數(shù)函數(shù)名要和 Relay 側 lowering 時查找的名字一致。通常的約定是topi.myop.leaky_clamp這種形式。Relay 在 lowering 時會根據(jù)算子的名字拼出這個全局函數(shù)名去調(diào)用。這里有個非常容易踩的坑compute 函數(shù)的參數(shù)順序。它必須和 Relay 側add_argument加屬性字段的順序完全一致否則運行時會因為參數(shù)錯位產(chǎn)生莫名其妙的錯誤。我建議在寫的時候把參數(shù)列表寫下來對照不要憑記憶。4.2 注冊 compute 與 lowering 對接compute 寫好后注冊代碼通常長這樣TVM_REGISTER_GLOBAL(topi.myop.leaky_clamp) .set_body([](TVMArgs args, TVMRetValue* rv) { *rv LeakyClampCompute(args[0], args[1], args[2], args[3]); });或者更常見的是在 Python 側用tvm.register_func注冊tvm.register_func(topi.myop.leaky_clamp) def _leaky_clamp_compute(data, lower, upper, slope): return leaky_clamp(data, lower, upper, slope)兩種方式都可以C 側注冊性能稍好Python 側注冊更靈活、改起來不用重新編譯。開發(fā)階段我強烈建議先用 Python 注冊因為改一行代碼就能生效不用等 C 編譯。等邏輯穩(wěn)定了再考慮挪到 C。lowering 對接的關鍵是 Relay 怎么知道去調(diào)用哪個全局函數(shù)。Relay 的OpImplementation或者FTVMCompute屬性會指定 compute 函數(shù)名。在RELAY_REGISTER_OP里可以用set_attrFTVMCompute(FTVMCompute, ...)來指定或者依賴默認的命名約定。默認約定是算子名加前綴比如myop.leaky_clamp對應topi.myop.leaky_clamp。如果你用了非默認命名一定要顯式設置FTVMCompute。4.3 編譯、測試與驗證代碼寫完之后編譯 TVM。如果你只改了 Python 側不需要重新編譯直接跑測試就行。如果改了 C 側需要重新編譯這個過程可能比較久建議用ninja增量編譯。測試分幾個層次。第一層是算子能不能被正確解析和類型推導寫個簡單的 Relay 程序用relay.Function構造調(diào)用跑relay.transform.InferType看輸出類型對不對。第二層是數(shù)值正確性構造隨機輸入用新算子算一遍再用 NumPy 手動實現(xiàn)算一遍對比結果。第三層是端到端把算子放進一個完整模型編譯到目標后端跑推理看結果。數(shù)值測試的代碼大概這樣import numpy as np import tvm from tvm import relay data relay.var(data, relay.TensorType((2, 3), float32)) out relay.op.myop.leaky_clamp(data, lower-1.0, upper1.0, slope0.01) func relay.Function([data], out) mod tvm.IRModule.from_expr(func) mod relay.transform.InferType()(mod) x np.random.randn(2, 3).astype(float32) expected np.maximum(np.minimum(x, 1.0), -1.0) 0.01 * np.minimum(x 1.0, 0.0) with tvm.transform.PassContext(opt_level3): lib relay.build(mod, targetllvm) runtime tvm.contrib.graph_executor.GraphModule(lib[default](tvm.cpu())) runtime.set_input(data, x) runtime.run() result runtime.get_output(0).numpy() np.testing.assert_allclose(result, expected, rtol1e-5)這個測試覆蓋了類型推導、lowering、編譯、執(zhí)行全流程。如果這一步過了說明算子基本可用。如果沒過錯誤信息通常會指向某個具體階段按階段排查即可。5. 常見問題與排查技巧實錄5.1 編譯期報錯注冊相關的典型問題新增算子時編譯期報錯主要集中在注冊環(huán)節(jié)。最常見的是“算子未注冊”或者“找不到類型關系”。前者通常是RELAY_REGISTER_OP的名字和調(diào)用時用的名字不一致注意 Relay 里算子名是大小寫敏感的而且命名空間前綴要完全匹配。后者是add_type_rel綁定的名字和實際 TypeRelation 注冊的名字不一致。另一個常見問題是屬性結構體沒注冊。如果你定義了LeakyClampAttrs但忘了用TVM_REGISTER_NODE_TYPE注冊編譯時可能不報錯但運行時解析屬性會失敗。注冊代碼是TVM_REGISTER_NODE_TYPE(LeakyClampAttrs);這個宏要放在屬性結構體定義之后。我踩過一次坑屬性結構體定義在頭文件里注冊宏忘了加結果算子能注冊成功但一調(diào)用就崩排查了半天才發(fā)現(xiàn)是屬性沒注冊。還有一種情況是set_num_inputs和實際add_argument的次數(shù)不匹配。set_num_inputs聲明的是數(shù)據(jù)輸入個數(shù)add_argument也要調(diào)用相同次數(shù)。如果聲明了 2 個輸入但只add_argument了 1 次類型推導時types數(shù)組長度會對不上報數(shù)組越界的錯。5.2 運行期報錯類型推導與 lowering 的坑運行期報錯里類型推導相關的最多。典型癥狀是Check failed: reporter-Assign之類的錯誤通常是因為輸出類型和輸入類型不兼容或者Assign被調(diào)用了多次。檢查你的 TypeRelation確保每個輸出只Assign一次且 dtype 和 shape 都合理。lowering 階段的報錯通常是“找不到 compute 函數(shù)”。這多半是全局函數(shù)名和 Relay 期望的名字不一致。排查方法是打印出 Relay 期望的函數(shù)名或者直接在RELAY_REGISTER_OP里顯式設置FTVMCompute屬性。另外compute 函數(shù)返回的必須是 TE 的 Tensor 或者 Tensor 列表如果返回了別的東西lowering 會失敗。數(shù)值不對但能跑通的情況通常是 compute 實現(xiàn)里的索引或者常量寫錯了。比如te.min和te.max的參數(shù)順序搞反或者常量 dtype 和輸入 dtype 不一致導致隱式類型轉換。這類問題建議先用小 shape 的輸入手動驗證比如(1, 1)或者(2, 2)把中間結果打印出來對比。5.3 常見問題速查表問題現(xiàn)象可能原因排查方向算子未注冊名字不匹配或注冊宏未執(zhí)行檢查RELAY_REGISTER_OP名字與調(diào)用名類型關系找不到add_type_rel名字與注冊名不一致核對 TypeRelation 注冊字符串屬性解析失敗屬性結構體未注冊補TVM_REGISTER_NODE_TYPE數(shù)組越界set_num_inputs與add_argument次數(shù)不符核對輸入個數(shù)聲明找不到 compute全局函數(shù)名不匹配顯式設置FTVMCompute或核對命名數(shù)值錯誤compute 索引或常量寫錯小 shape 手動驗證中間結果融合失效tag 設置不當檢查 compute 的 tag 參數(shù)5.4 獨家避坑經(jīng)驗第一條經(jīng)驗開發(fā)階段全部用 Python 注冊 compute不要碰 C。C 編譯一次動輒十幾分鐘改一個符號就要重編效率極低。等 Python 版本驗證通過再考慮要不要挪到 C。實際上很多自定義算子用 Python 注冊完全夠用性能差異在大多數(shù)場景下可以忽略。第二條經(jīng)驗類型關系先用現(xiàn)成的實在不行再自己寫。ElemwiseRel和BroadcastRel能覆蓋大部分逐元素和廣播場景自己寫 TypeRelation 容易在邊界情況上出錯。如果非寫不可先寫一個最簡版本只處理TensorType把nullptr和動態(tài) shape 的情況留到后面再補。第三條經(jīng)驗測試要從最小用例開始。不要一上來就用真實模型的復雜輸入先用(1, 1)或者(2, 3)這種小 shape把類型推導、lowering、數(shù)值都對一遍。小用例跑通了再逐步放大。我見過有人直接上大模型報錯信息淹沒在幾百行日志里根本定位不到問題。第四條經(jīng)驗善用relay.transform.InferType和relay.transform.PrintIR。類型推導出問題時先跑InferType看報錯位置lowering 出問題時用PrintIR把 lowering 前后的 IR 打出來對比能快速定位是哪一步出的問題。第五條經(jīng)驗屬性字段的默認值要合理。如果某個屬性沒有默認值用戶調(diào)用時必須顯式傳否則解析會失敗。給一個合理的默認值能降低使用門檻也能避免因為漏傳參數(shù)導致的詭異錯誤。6. 算子擴展后的優(yōu)化與融合考量6.1 讓新算子參與算子融合新算子加進去只是第一步能不能融入 Relay 的優(yōu)化流程才是決定性能的關鍵。Relay 的融合 pass 會把相鄰的算子合并成一個 kernel減少內(nèi)存訪問和 kernel 啟動開銷。新算子能不能被融合取決于它的 tag 和融合模式注冊。compute 里的tag參數(shù)是融合 pass 的重要依據(jù)。elemwise標簽的算子通常可以和前后算子自由融合broadcast次之reduce和injective有各自的融合規(guī)則。如果你希望新算子能融合tag 要設對。比如逐元素算子設elemwise涉及廣播的設broadcast。除了 tag還可以用register_fusion_pattern注冊更精細的融合模式。比如你希望某個算子和它前面的卷積融合可以寫一個模式匹配規(guī)則。這部分相對進階不是所有算子都需要但如果你的算子性能敏感值得花時間研究。6.2 自定義 schedule 的接入點如果新算子的 compute 用默認 schedule 性能不理想可以自定義 schedule。TOPI 的 compute 函數(shù)通常配套一個 schedule 函數(shù)用te.create_schedule創(chuàng)建調(diào)度然后對循環(huán)做切分、向量化、并行化等操作。自定義 schedule 的接入方式是在 Relay 的 lowering 配置里指定。TVM 的relay.build可以傳schedule_cfg或者用OpStrategy機制把 compute 和 schedule 綁定。對于自定義算子最簡單的做法是在 TOPI 里提供一個 schedule 函數(shù)然后在 lowering 時通過FTVMSchedule屬性指定。schedule 的編寫是另一個大話題這里不展開。但要提醒一點不要過早優(yōu)化 schedule。先用默認 schedule 把功能跑通測出性能瓶頸再針對性地寫 schedule。很多時候默認 schedule 已經(jīng)夠用自定義 schedule 的收益可能還不如換個更合適的 tag 讓融合 pass 發(fā)揮作用。6.3 跨后端兼容性檢查新算子寫完之后如果你需要它在多個后端上都能跑要做兼容性檢查。不同后端對算子的支持程度不一樣比如某些后端不支持特定的數(shù)據(jù)類型或者不支持動態(tài) shape。測試的時候至少覆蓋 CPU 和你有需求的目標后端。如果某個后端不支持通常有兩種處理方式。一種是提供 fallback 實現(xiàn)讓不支持的部分回退到 CPU 或者用其他算子組合。另一種是在 lowering 階段做后端特定的轉換把新算子拆解成后端支持的算子。具體選哪種取決于你的性能要求和實現(xiàn)成本。還有一個容易被忽略的點是 dtype 支持。你的 compute 實現(xiàn)里如果用了tvm.tir.const包裝常量要確保常量的 dtype 和輸入 dtype 一致。如果輸入可能是float16也可能是float32compute 要能處理兩種 dtype不能寫死。測試的時候把常見 dtype 都跑一遍避免上線后才發(fā)現(xiàn)某個 dtype 不支持。7. 我個人在實際操作中的幾點體會走完整個流程我最深的體會是TVM 的算子注冊機制雖然看起來繁瑣但它的設計是有內(nèi)在邏輯的。注冊、類型關系、compute 這三層分別對應“是什么”“怎么推”“怎么算”把這三層想清楚代碼就是水到渠成的事。真正難的不是寫代碼而是理解每一層為什么存在、什么時候需要自定義、什么時候可以復用。另一個體會是文檔和社區(qū)示例很重要但不能全信。TVM 版本迭代快不同版本的 API 可能有差異網(wǎng)上找到的示例代碼不一定適用于你的版本。最可靠的做法是直接看當前版本 TVM 源碼里已有的算子實現(xiàn)照葫蘆畫瓢。比如你想加一個逐元素算子就找一個現(xiàn)有的逐元素算子把它的注冊、類型關系、compute 都讀一遍然后照著改。這比看任何教程都管用。最后一點不要怕報錯。新增算子涉及的環(huán)節(jié)多第一次做幾乎不可能一次通過。報錯信息通常很具體指向某個文件某一行順著排查基本都能找到原因。我自己的習慣是每改一小步就編譯測試一次不要攢一大堆改動再一起測那樣出了問題很難定位。小步快跑逐步驗證是走通這個流程最穩(wěn)的方式。如果你后續(xù)還想擴展可以考慮的方向包括給算子加梯度支持讓它能參與訓練、寫自定義 schedule 提升性能、或者把算子貢獻回 TVM 社區(qū)。每個方向都有各自的坑但核心流程和這篇文章講的是一致的。把基礎流程走通后面的擴展就是在這個框架上做加法。