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