1. 一个 shape 对不上的下午把 IR、计算类型、dtype、infer shape 串成一条线前阵子帮一位做端侧部署的朋友排查一个小模型输入是(1, 3, 224, 224)的float32图像最后卡在了一句非常没有营养的报错上类型推导出来的张量形状和下游算子期望的对不上堆栈往上一翻全是编译器的内部 pass 名一行用户代码都看不到。他一开始以为是模型转换的问题把 ONNX 转来转去转了三遍我让他先把 IR 打出来看看结果发现是某个自定义后处理算子里的te.compute写错了 shape 维度——用了一个两维的输入却只在 lambda 里取了一维的索引编译期没报错到 lowering 之后才炸。这类问题在 TVM 这套深度学习编译栈里太常见了。原因不在于难度而在于它把一件事拆成了四层IR 定义决定了你用什么样的节点描述计算计算类型表达式节点类型与算子 pattern决定了编译器能不能对这段计算做融合和调度dtype决定了数值在内存里怎么排布、运算怎么提升infer shape则是在没有真实数据的前提下把每一层的张量形状和数据类型推出来。这四层任何一层没对齐报错信息就会漂移到很远的地方看起来跟你的改动毫无关系。所以这篇东西我想换个顺序讲不从 API 手册式的目录往下铺而是从这四个概念之间到底怎么咬合入手把 Relay IR 和 TIR 各自负责什么、计算类型到底指哪几种东西、DLDataType那套编码规则怎么读、shape 和 type 是怎么一步步被推导出来的全部拆开讲清楚。适合已经能跑通 TVM 的 hello world、但在写自定义算子或调 pass 时经常被莫名其妙的类型错误挡住的人也适合只是想搞清楚编译器内部大概在干什么、不想深挖调度的读者。看完之后你至少能做到两件事看到 IR 打印结果能读懂它在说什么遇到 shape 或 dtype 报错知道该去哪个环节找。需要先提一句版本背景这个很影响你照着敲代码的结果。0.10到0.12这一段高层图 IR 是 RelayTensorIR 已经能用T.prim_func写了0.13之后新的 Relax 体系开始引入Relay 还在但不再是唯一选择。所以你敲代码之前先print(tvm.__version__)确认一下别拿着新版文档在老版本上死磕反过来也一样。2. Relay 与 TIR两种抽象层级的 IR 各自管什么很多人第一次看 TVM 的文档会晕因为它同时存在两套看起来都能描述计算的 IR而且同一个算子往往要写两遍。这不是设计冗余而是两个不同的抽象层级Relay 描述的是算什么TIR 描述的是怎么算。2.1 Relay IR 的骨架Expr 是节点Module 是容器Relay 是一套纯函数式、带类型的高层图表示。它的基本单位是Expr所有节点都继承自它。你日常会碰到的就那么几个Constant编译期常量relay.const(1)默认是int32这一点后面讲 dtype 时还会踩到。Var图的输入或中间变量带shape和dtype两个注解。Call对Op的调用是 Relay 里最重的节点算子的实际计算全挂在它身上。Tuple/TupleGetItem多输出算子的表达方式卷积 池化这类融合前经常会遇到。Let/If/Function控制流和函数抽象Function是IRModule里main那个东西的类型。把这些节点装进一个IRModule用GlobalVar指向每个Function就是一份完整的 Relay 程序了。这份程序是不可变的——你想改它就得用 pass 生成一份新的IRModule理解这一点对后面调InferType很关键因为它不是就地修改你得接住返回值。from tvm import relay import tvm x relay.var(x, shape(1, 3, 224, 224), dtypefloat32) w relay.var(w, shape(64, 3, 7, 7), dtypefloat32) y relay.nn.conv2d(x, w, padding(3, 3), kernel_size(7, 7), channels64) mod tvm.IRModule.from_expr(y) mod relay.transform.InferType()(mod) print(mod[main].body.checked_type)跑完这段checked_type里会是一个TensorType形状和 dtype 都在里面。注意checked_type这个属性只有跑过InferType的表达式才有直接对原始expr取会报错。这是我见过新手问得最多的一个问题之一。2.2 TIR落到 buffer、循环和线程绑定Relay 层看不到内存也看不到循环。真正能生成代码的是 TIR。它的核心对象是PrimFunc主体由Buffer、For带ForKind、BufferLoad、BufferStore、IfThenElse、SeqStmt、Evaluate这些节点组成。ForKind是个容易被忽略但很有意思的设计它直接编码了这个循环属于什么计算类型ForKind含义典型来源kSerial普通串行循环te.compute默认产物kUnrolled强制展开schedule.unrollkVectorized向量化配合lanes 1的 dtypeschedule.vectorizekParallel多线程并行schedule.parallelkThreadBinding绑定到 GPU 线程轴bindthread_axis用 TVMScript 写一个最小PrimFunc比手搓ir_builder舒服得多import tvm from tvm.script import tir as T T.prim_func def scale(A: T.Buffer((128,), float32), B: T.Buffer((128,), float32)): for i in T.serial(128): B[i] A[i] * 2.0这段代码里128是编译期固定的。如果换成符号维度就得用T.int32类型的Var或者te.var写法会变后面第 5 节讲动态 shape 时会具体展开。2.3 两者的交界处最容易出错的地方Relay 和 TIR 之间靠FTVMCompute这个属性连起来。一个 Relay op 注册的时候除了告诉编译器我有几个输入、类型关系是什么还要提供一个计算函数把Arrayte::Tensor转成 TIR 能被 lower 的东西。如果你写自定义算子最常见的错误就出在这里TypeRelation 里推导出来的 shape 和你 compute 函数里实际产出的 shape 对不上。编译器在融合之前会先做类型推导拿到TensorType然后又拿这个TensorType去驱动 compute 生成 TE 张量。这两边只要有一个维度不一样报错就会出现在很远的地方——可能是 schedule 阶段说某个 axis 越界也可能是 lower 之后生成的循环访问越界。我一般的做法是写完 compute 之后立刻print(tvm.lower(s, [A, B], simple_modeTrue))把 TIR 打出来肉眼核对循环边界跟你想的是不是一回事。这一步花不了十秒钟能省掉后面半小时的猜测。还有一个隐性的边界是 layout。Relay 层的nn.conv2d默认是 NCHW但后端可能要求 NHWC这个转换由FInferCorrectLayout这类属性驱动。你自己写的算子如果没注册这个属性融合的时候就可能拿到一个布局和你预期不一致的输入然后在 compute 里按错误的下标取值——这类 bug 不会报错只会算错最难受。3. 计算类型从 Expr 节点家族到算子的 pattern 归类计算类型这个词在 TVM 语境里其实指向两个不同的东西我把它俩分开讲因为混在一起会让人非常困惑一个是表达式节点的类型Expr的继承体系另一个是算子的计算模式标签OpPatternKind。前者是结构概念后者是调度概念。3.1 Expr 节点类型决定了你能对 IR 做什么Expr的继承体系决定了 pass 能怎么处理一段程序。举个具体例子Call节点里挂了Op、参数列表、Attrs、类型参数四个字段。你想做一个算子替换的 pass本质就是递归遍历Expr遇到Call检查它的Op是不是目标算子是就构造新的Call换掉。如果你不知道Tuple和TupleGetItem的存在遍历到多输出算子时就会漏掉pass 跑完看起来成功结果图里还有一半没换。Op本身也是个节点类型你可以用tvm.ir.Op.get(nn.conv2d)把它取出来查它注册了哪些属性。这个操作在排查为什么某个算子没有被融合时特别好用op tvm.ir.Op.get(nn.conv2d) print(op.num_inputs) print(op.get_attr(TOpPattern)) # 返回 OpPatternKind 的数值 print(op.get_attr(TOpIsStateful))get_attr返回的是ObjectRef或者数字TOpPattern那个会返回一个整数对应OpPatternKind枚举。你如果只打印数字看不懂可以对照下一小节的含义表。3.2 OpPattern 是给调度器和算子融合看的标签OpPatternKind是 Relay 层给算子贴的一个我能被怎么优化的标签。融合 passFuseOps遍历图的时候就是看相邻节点的 pattern 决定能不能合成一个Function。常见取值和含义Pattern含义融合行为kElemWise逐元素输入输出形状一致可以自由融合进上游kBroadcast带广播的逐元素同 elemwise但要处理 shape 不同kInjective输入输出元素一一对应但形状可变如 reshape、transpose可融合一般不增加计算量kCommReduce归约输出维度少于输入通常作为融合的边界kOutEWiseFusable输出侧可融合如 conv作为融合的锚点kOpaque不透明无法分析阻断融合标签贴错后果很直接。我见过有人把自定义的归一化算子标成kElemWise但它内部实际上做了跨通道归约。融合 pass 一看是 elemwise毫不犹豫地和前后的算子合成一个 kernel结果调度器按逐元素的方式生成循环归约那个维度被当成了普通循环计算结果全乱。判断标准很简单输入输出的形状和元素是否一一对应。只要有一个元素依赖其他位置的元素就不该是 elemwise。Python 侧改这个标签不需要重新编译整个 TVM用注册接口就行tvm.ir.register_op_attr(my_scale, TOpPattern) def _(op): return tvm.relay.op.OpPatternElemWise注意这是注册不是修改同一个 key 重复注册会报重复定义调试的时候记得把旧的注册代码注释掉不然你会怀疑人生。3.3 从 pattern 反推算子该怎么写我自己的习惯是先想清楚这个算子的 pattern再决定 IR 怎么写。如果是kElemWisete.compute的 shape 直接照抄输入 shapelambda 里只用当前索引这种写法几乎不会错。如果是kCommReduce就必须有reduce_axisshape 要比输入少一维而且你要注意te.sum这类归约原语的axis参数必须是reduce_axis对象不能随手传个整数。from tvm import te m, k 128, 256 A te.placeholder((m, k), nameA, dtypefloat32) kk te.reduce_axis((0, k), namekk) B te.compute((m,), lambda i: te.sum(A[i, kk], axiskk), nameB)这段里i是自由变量kk是被归约掉的轴B的 shape 只有(m,)。如果你不小心把axiskk写成axisiTVM 会直接报错说 axis 不在输入维度范围内这个报错还算友好。真正难查的是 shape 里少了维度但恰好能广播的情况算出来是错的但不报错。顺带说一个 C 侧注册算子的骨架方便你对照理解这三个东西的对应关系RELAY_REGISTER_OP(my_scale) .set_num_inputs(1) .add_argument(data, Tensor, 输入张量) .set_support_level(1) .add_type_relMyScaleRel(MyScale) // 管 shape 和 dtype .set_attrTOpPattern(TOpPattern, kElemWise) // 管融合 .set_attrFTVMCompute(FTVMCompute, MyScaleCompute); // 管实际计算三行属性各司其职缺一行都有对应的症状缺 type rel 会在InferType阶段挂掉缺TOpPattern会导致融合 pass 把它当成 opaque 直接切断缺FTVMCompute会在 build 阶段告诉你找不到计算函数。4. dtype四位十六进制背后的一套编码规则dtype 看起来只是个字符串float32但 TVM 内部把它压成了一个紧凑的位域结构。理解这个结构能解释掉一大批看起来莫名其妙的报错。4.1 DLDataType 的位域与 lanes 的含义TVM 的DataType底层对应三字段code类型族、bits位宽、lanes通道数。code的取值沿用了一套跨框架通用的约定大致是0表示有符号整型、1表示无符号整型、2表示浮点、3表示不透明句柄、4表示 bfloat、5表示复数、6表示布尔。新版还补了若干 8 位浮点的编码。你在 Python 里可以直接拆开看dt tvm.DataType(float32) print(dt.bits, dt.type_code, dt.lanes) v tvm.DataType(int8x4) print(v.bits, v.lanes) # 8, 4lanes是最容易被忽略的一位。float32的lanes是 1int8x4的lanes是 4表示这是一个由 4 个 8 位整数组成的向量类型。向量化的时候schedule 会把循环的ForKind改成kVectorized然后 TIR 里的 buffer 访问就会用上带 lanes 的 dtype。这也是为什么你在写自定义算子的时候如果硬编码了float32而用户传了float32x4会在类型检查阶段直接失败。dtype的合法性判断可以提前做省得编译到一半才炸assert tvm.DataType(float16).lanes 1 # 非法写法会直接抛异常 try: tvm.DataType(float33) except Exception as e: print(非法 dtype:, e)4.2 Relay 的类型推导链路与 TensorTypeRelay 里 shape 和 dtype 是绑在一起推导的载体是TensorType。这个类型有两个字段shape一个ArrayPrimExpr和dtype一个DataType。InferType这个 pass 干的事就是自底向上遍历Expr对每个节点调用它注册的 type relation把输出类型填进checked_type_。对二元运算来说type relation 通常是BroadcastRel它会做两件事检查两个输入的 shape 能不能广播、检查两个输入的 dtype 是否完全一致。第二点特别反直觉——Relay 不做隐式类型提升。你写惯了 NumPy会觉得int32 float32应该自动变成float64或者float32但在 Relay 里这是错误会直接报类型不兼容。x relay.var(x, shape(2, 3), dtypefloat32) y relay.var(y, shape(2, 3), dtypeint32) bad relay.add(x, y) # 构造不会报错 mod tvm.IRModule.from_expr(bad) try: mod relay.transform.InferType()(mod) # 这里才炸 except Exception as e: print(类型不匹配:, e)注意报错时机构造relay.add的时候是不检查的只有InferType才会检查。所以如果你的 pass 在InferType之前就跑可能会把错误传播得更远堆栈更难读。我的习惯是在自定义 pass 前后各插一次InferType让问题尽早暴露。4.3 混合精度和整型提升里最常见的三个坑第一个坑是常量 dtype。relay.const(1)的默认 dtype 是int32relay.const(1.0)是float32。很多人想给float16的输入加个常数顺手写了relay.const(1)然后在InferType阶段一脸茫然。正确写法是显式指定relay.const(1, float16)。第二个坑是标量与张量的广播。Relay 的二元运算对零维张量的处理和你想的可能不一样稳妥的做法是先用relay.op.broadcast_to_like展开成同 shape 再算或者干脆在 TE 层用标量常数——那边te.compute里的 Python 浮点字面量会自动转成匹配的 dtype省事很多。第三个坑是量化相关的整型累加。做 int8 量化推理时卷积的累加结果通常是 int32如果中间某一步不小心用了 int8 buffer 存储部分和会静默溢出数值全部错乱但不报错。这类问题唯一的解法就是在 TIR 层把 buffer 的 dtype 打出来逐个核对别信应该没问题。5. infer shape 的三种实现路径与实际写法shape 推导在不同层有三套不同的机制这个分法我觉得比infer shape 是个 pass这种笼统说法有用得多因为它决定了你遇到问题时该去哪儿改。5.1 te.compute 的 shape lambda最直接也最容易写错TE 层的 shape 是你直接给的不需要推。te.compute(shape, f)里的shape参数就是输出张量的形状f的返回值是每个位置的元素表达式。这里最容易错的地方是**shape的长度必须和 lambda 的参数个数一致**。lambda i, j: ...对应两维 shape少写一个参数或者多写一个报错信息都不太直观。另一个高频错误是符号维度的来源不一致。你用一个te.var(n)定义了输入的 shape然后输出 shape 里随手写了 Python 的int这在静态 shape 下没问题但一旦n是符号的生成的 TIR 就会变成两个互不相关的维度lower 之后循环边界对不上。n te.var(n) A te.placeholder((n,), nameA, dtypefloat32) B te.compute((n,), lambda i: A[i] 1.0, nameB) # 正确用同一个 n C te.compute((128,), lambda i: A[i] 1.0, nameC) # 危险n 和 128 没关系C这种写法在n ! 128时会在 lower 阶段报越界或者生成错误代码。所以我的原则是只要有一个维度是符号的所有相关的 shape 表达式都必须从同一个符号源头派生用n * 2、n // 4这类表达式而不是新的独立变量。5.2 Relay op 的 TypeRelationshape 与 type 一起推Relay 层的推导走的是TypeRelation你需要实现一个类似下面这样的结构把输入类型读出来构造输出类型然后Assign给输出位置struct MyScaleRel : public TypeRelation { bool TypeRel(const ArrayType types, int num_inputs, const Attrs attrs, const TypeReporter reporter) { ICHECK_EQ(types.size(), 2U); const auto* data types[0].asTensorTypeNode(); if (data nullptr) return false; // 输入类型还没推出来先放弃 reporter-Assign(types[1], TensorType(data-shape,>x relay.var(x, shape(relay.Any(), 3, 224, 224), dtypefloat32)但要注意Any只是让类型推导别报错它并不会让下游算子自动适配。很多算子在 shape 里出现Any时type relation 会直接返回 false整个InferType就卡住了。这时候通常的做法是先DynamicToStatic把已知的 batch 固定下来跑完优化之后再切回去或者用 shape 函数relay.Function形式在运行时把 shape 算出来。TIR 侧的动态 shape 用的是tvm.tir.Any()或者符号Var两者语义不同Var表示这是个会传进来的值Any表示真的不知道。Any在 TIR 里几乎没法用来生成循环因为它不是一个可计算的表达式只能出现在 buffer 的 shape 声明里做占位。真正要生成代码还是得把它替换成Var。T.prim_func def dyn(A: T.Buffer((T.int32(n),), float32), B: T.Buffer((T.int32(n),), float32)): for i in T.serial(T.int32(n)): B[i] A[i] * 2.0这种写法的关键点是任务调度时要把n作为参数传进去否则生成的 kernel 不知道循环边界。这也是为什么端侧部署时大家更偏爱固定 batch——动态 shape 带来的不只是性能开销还有一整条链路上每个环节都要处理符号变量的麻烦。6. 我平时排查 shape 与 dtype 问题的固定套路把上面这些东西揉在一起落到实际操作上我有一套基本固定的检查顺序。分享出来你照着走一遍大部分问题能定位得比读堆栈快。第一步永远先打 IR。print(mod.astext(show_meta_dataFalse))这一句能解决 30% 的困惑因为很多问题的根源是你以为图长这样实际上不是。尤其是经过几轮 pass 之后图可能已经被改得面目全非了。我一般会在每个关键 pass 前后各打一次虽然输出很长但用grep过滤算子名很快。第二步强制跑一次InferType并打印checked_type。这一步能确认 shape 和 dtype 到底在哪一层开始偏离预期。如果InferType就报错问题在算子注册的 type relation 或者输入变量的注解如果InferType通过但 build 失败问题在 compute 函数或者 schedule。第三步lower 成 TIR 看循环边界。tvm.lower(s, args, simple_modeTrue)的输出可读性很好能直接看到每个循环的起止值、buffer 的访问下标。shape 推导错了的话这里一定会暴露成越界或者维度缺失。第四步对着报错里的算子名用tvm.ir.Op.get(name).get_attr(...)逐个查它注册了什么属性。缺失的属性能直接指出问题类型没有TOpPattern说明融合行为不可控没有FTVMCompute说明 build 阶段必挂type relation 返回 false 说明推导链断了。症状大概率原因检查点InferType报 not compatible两侧 dtype 不一致常量 dtype、输入 var 注解InferType卡住不返回type relation 一直返回 false输入类型是否可推出、是否用了 Anybuild 阶段找不到计算函数缺FTVMCompute算子注册代码编译通过但结果错pattern 标错或 layout 假设错TOpPattern、FInferCorrectLayoutlower 后循环边界为 0shape 表达式用了独立符号符号变量是否同源最后再补一个我自己踩过的坑说出来可能有点丢人早期我写自定义算子习惯性地把te.compute的 shape 写成输入 shape 的 Python 拷贝也就是inputs[0].shape。在静态 shape 下这完全没问题但只要输入是符号的inputs[0].shape里的Var和你后面新建的Var是两个不同的对象虽然名字可能都叫n编译器会认为是两个独立维度。判断方法是打印tvm.ir.structural_equal(a, b)或者干脆看 lower 出来的 TIR 里循环边界是不是同一个变量名——同名但不同对象的符号在打印时长得一模一样这是最坑的地方。后来我养成一个习惯所有符号 shape 都从inputs[0].shape里取绝不自己新建这样能规避掉一大类看起来对但就是不对的问题。