企业官网建设流程全解析

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 是 Relay，TensorIR 已经能用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), dtype="float32") w = relay.var("w", shape=(64, 3, 7, 7), dtype="float32") y = relay.nn.conv2d(x, w, padding=(3, 3), kernel_size=(7, 7), channels=64) 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是个容易被忽略但很有意思的设计，它直接编码了"这个循环属于什么计算类型"：

用 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 注册的时候，除了告诉编译器"我有几个输入、类型关系是什么"，还要提供一个计算函数，把Array<te::Tensor>转成 TIR 能被 lower 的东西。如果你写自定义算子，最常见的错误就出在这里：TypeRelation 里推导出来的 shape 和你 compute 函数里实际产出的 shape 对不上。

编译器在融合之前会先做类型推导，拿到TensorType；然后又拿这个TensorType去驱动 compute 生成 TE 张量。这两边只要有一个维度不一样，报错就会出现在很远的地方——可能是 schedule 阶段说某个 axis 越界，也可能是 lower 之后生成的循环访问越界。我一般的做法是：写完 compute 之后立刻print(tvm.lower(s, [A, B], simple_mode=True))把 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 层给算子贴的一个"我能被怎么优化"的标签。融合 pass（FuseOps）遍历图的时候，就是看相邻节点的 pattern 决定能不能合成一个Function。常见取值和含义：

标签贴错后果很直接。我见过有人把自定义的归一化算子标成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 怎么写。如果是kElemWise，te.compute的 shape 直接照抄输入 shape，lambda 里只用当前索引，这种写法几乎不会错。如果是kCommReduce，就必须有reduce_axis，shape 要比输入少一维，而且你要注意te.sum这类归约原语的axis参数必须是reduce_axis对象，不能随手传个整数。

from tvm import te m, k = 128, 256 A = te.placeholder((m, k), name="A", dtype="float32") kk = te.reduce_axis((0, k), name="kk") B = te.compute((m,), lambda i: te.sum(A[i, kk], axis=kk), name="B")

这段里i是自由变量，kk是被归约掉的轴，B的 shape 只有(m,)。如果你不小心把axis=kk写成axis=i，TVM 会直接报错说 axis 不在输入维度范围内，这个报错还算友好。真正难查的是 shape 里少了维度但恰好能广播的情况，算出来是错的但不报错。

顺带说一个 C++ 侧注册算子的骨架，方便你对照理解这三个东西的对应关系：

RELAY_REGISTER_OP("my_scale") .set_num_inputs(1) .add_argument("data", "Tensor", "输入张量") .set_support_level(1) .add_type_rel<MyScaleRel>("MyScale") // 管 shape 和 dtype .set_attr<TOpPattern>("TOpPattern", kElemWise) // 管融合 .set_attr<FTVMCompute>("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, 4

lanes是最容易被忽略的一位。float32的lanes是 1，int8x4的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 的类型推导链路与 TensorType

Relay 里 shape 和 dtype 是绑在一起推导的，载体是TensorType。这个类型有两个字段：shape（一个Array<PrimExpr>）和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), dtype="float32") y = relay.var("y", shape=(2, 3), dtype="int32") 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 是int32，relay.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,), name="A", dtype="float32") B = te.compute((n,), lambda i: A[i] + 1.0, name="B") # 正确：用同一个 n C = te.compute((128,), lambda i: A[i] + 1.0, name="C") # 危险：n 和 128 没关系

C这种写法在n != 128时会在 lower 阶段报越界或者生成错误代码。所以我的原则是：只要有一个维度是符号的，所有相关的 shape 表达式都必须从同一个符号源头派生，用n * 2、n // 4这类表达式而不是新的独立变量。

5.2 Relay op 的 TypeRelation：shape 与 type 一起推

Relay 层的推导走的是TypeRelation，你需要实现一个类似下面这样的结构，把输入类型读出来，构造输出类型，然后Assign给输出位置：

struct MyScaleRel : public TypeRelation { bool TypeRel(const Array<Type>& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { ICHECK_EQ(types.size(), 2U); const auto* data = types[0].as<TensorTypeNode>(); if (data == nullptr) return false; // 输入类型还没推出来，先放弃 reporter->Assign(types[1], TensorType(data->shape,>x = relay.var("x", shape=(relay.Any(), 3, 224, 224), dtype="float32")

但要注意，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_data=False))这一句能解决 30% 的困惑，因为很多问题的根源是"你以为图长这样，实际上不是"。尤其是经过几轮 pass 之后，图可能已经被改得面目全非了。我一般会在每个关键 pass 前后各打一次，虽然输出很长，但用grep过滤算子名很快。

第二步，强制跑一次InferType并打印checked_type。这一步能确认 shape 和 dtype 到底在哪一层开始偏离预期。如果InferType就报错，问题在算子注册的 type relation 或者输入变量的注解；如果InferType通过但 build 失败，问题在 compute 函数或者 schedule。

第三步，lower 成 TIR 看循环边界。tvm.lower(s, args, simple_mode=True)的输出可读性很好，能直接看到每个循环的起止值、buffer 的访问下标。shape 推导错了的话，这里一定会暴露成越界或者维度缺失。

第四步，对着报错里的算子名，用tvm.ir.Op.get(name).get_attr(...)逐个查它注册了什么属性。缺失的属性能直接指出问题类型：没有TOpPattern说明融合行为不可控，没有FTVMCompute说明 build 阶段必挂，type relation 返回 false 说明推导链断了。

最后再补一个我自己踩过的坑，说出来可能有点丢人：早期我写自定义算子，习惯性地把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里取，绝不自己新建，这样能规避掉一大类"看起来对但就是不对"的问题。