TVM TE InferBound
https://tvm.apache.org/docs/v0.16.0/arch/inferbound.html
TVM TE InferBound Pass
InferBound Pass过程的主要任务是创建 bound 映射 rmap，指定每个 IterVar 的 Range。bound 将传递给 ScheduleOps，从而设置循环的范围，以及 buffer 的大小。
一个 TVM schedule 由多个 Stage 组成，每个 stage 实际上就是一个 Operation（如 ComputeOp 或 TensorComputeOp）。每个 operation 都有一个 root_iter_vars 的 IterVar 列表，例如，ComputeOp 的 root_iter_vars 由 axis 和 reduce axis 对应的 IterVar 组成。除此之外，每个 operation 也可以包含其他的 IterVars，但是这些 IterVars 是与 IterVarRelations （如 split、fuse 或 rebase）相关的。如 split 的 IterVarRelation，就指定了 split 操作的父 IterVar，以及两个孩子 inner 和 outer。
schedule 的多个 stage 组成了有向无环图（DAG），图中的每个节点即一个 stage。在 TVM 中，Tensor 表示 Operation 的输出。若 Stage B 的输入 Tensor 来源于 Stage A 的 op，那么 DAG 图中存在 Stage A 到 Stage B 的一条边。
InferBound 从输出节点开始，按反向拓扑序遍历 DAG 中的节点。InferBound 有两个属性：
- 访问一个节点 stage 后，rmap 中应该设置了该节点所有的 InterVar 的 range 映射。
- 在 rmap 中，每个 InterVar 的 Range 只会被设置一次。
在访问每个 stage 节点时，会调用两个函数，InferRootBound 和 PassDownDomain。InferRootBound 负责设置 root_iter_vars 中 IterVar 的 bound 映射；而 PassDownDomain 设置该 stage 其余 IterVar 的映射。
IterVar Hyper-graph
对于每个 Stage 节点，其内部结构包含一个以 InterVar 作为节点的有向无环超图，其超边是 IterVarRelation。InferRootBound 和 PassDownDomain 需要在这些 InterVar 图上进行消息传递。下图中，从父 InterVar 传向子 InterVar 的消息叫做 "PassDown" ，而反过来叫做 "PassUp"。
flowchart LR
classDef root fill:orange
classDef leaf fill:lightblue
classDef other fill:white
i([i]):::root --> Node0[SplitNode]
Node0 --> i.outer([i.outer]):::leaf
Node0 --> i.inner([i.inner]):::other
i.inner --> Node1[SplitNode]
Node1 --> i.inner.outer([i.inner.outer]):::leaf
Node1 --> i.inner.inner([i.inner.inner]):::leafPassDownDomain
InterVarRelation 有三种：split、fuse 以及 rebase。
-
split 的 inner 和 outer 的 range 需要通过父级节点的 range 进行计算：
-
fuse 的 range 需要通过其 inner 和 outer 的 range 来计算：
InferRootBound
InferRootBound 仅设置 root_iter_var 的 range。若该 stage 是一个输出 stage 或者是 placeholder，直接从 InterVar 的 dom 成员变量中获取 range。否则，需要进行 4 个阶段的操作。
Phase 1：初始化 consumer leaf_iter_vars 的 IntSet
- 在 InferRootBound 时，Range 需要被转换为 IntSet，即整数集合。
Phase 2：将 IntSet 从 consumer leaves 传播到 consumer 的 root
- 首先调用 PassUpDomain。
- 创建 root_iter_vars 到 IntSet 的映射
注意：若 schedule 中不包含 compute_at，那么 Phase 1-2 并不是必要的。
Phase 3：将 IntSet 传播到 consumer 的输入 tenser
consumer 的输入 tensor 是正在处理的 stage 的输出张量。对于张量的每个轴
Phase 4：整合所有 consumer
该阶段由 GatherBound 完成，其行为取决于 operation 的类型。对于 ComputeOp，它只有一个输出 tensor，ComputeOp 的 root_iter_vars 包含 axis 和 reduce_axis 变量。GatherBound 将 axis var 的 root InterVar 的 Range 设置为 consumer 对应 axis 的所有 IntSet 的并集（最小和最大值间的范围，这可能导致额外的计算）。而 reduce_axis 的 Range 设置为默认值（dom）。
compute_at 的 InferBound
若调度中包含 compute_at，那么 InferRootBound 的 Phase 1-2 会变得很复杂。
Motivation
例 1
C = tvm.compute((5, 16), lambda i, j : tvm.const(5, "int32"), name='C')
D = tvm.compute((5, 16), lambda i, j : C[i, j]*2, name='D')
stage D 需要 C 的所有 5*16 的元素用于计算。
例 2
C = tvm.compute((5, 16), lambda i, j : tvm.const(5, "int32"), name='C')
D = tvm.compute((5, 16), lambda i, j : C[i, j]*2, name='D')
s = tvm.create_schedule(D.op)
s[C].compute_at(s[D], D.op.axis[1])
仅需要 C 的一个元素
例 3
C = tvm.compute((5, 16), lambda i, j : tvm.const(5, "int32"), name='C')
D = tvm.compute((5, 16), lambda i, j : C[i, j]*2, name='D')
s = tvm.create_schedule(D.op)
s[C].compute_at(s[D], D.op.axis[0])
需要 C 的 16 个元素。
Attach Path
若 Stage C 在 Stage D 的轴 j 上 compute_at，那么说 C 被 attach 到 D 的 j 轴。为了确认 C 必须要计算的元素数量，需要了解 C 的计算是发生在 D 的 leaf 变量的哪个位置。CreateAttachPath 负责找到哪些 scope 包含 stage C，这些 scope 从内到外按顺序排列。例 1 中，C 的 attach path 是空的；例 2 中，是 {j,i}；例 3 中，是 {i}。
例 4
C = tvm.compute((5, 16), lambda i, j : tvm.const(5, "int32"), name='C')
D = tvm.compute((4, 5, 16), lambda di, dj, dk : C[dj, dk]*2, name='D')
s = tvm.create_schedule(D.op)
s[C].compute_at(s[D], D.op.axis[2])
这个例子中，C 的 attach path 是 {dk, dj, di}。IR 如下所示：
for di 0, 4
for dj 0, 5
for dk 0, 16
for i 0, 1
for j 0, 1
C[i+dj, j+dk] = 5
D[di, dj, dk] = C[dj, dk] * 2
例 5
C = tvm.compute((5, 16), lambda i, j : tvm.const(5, "int32"), name='C')
D = tvm.compute((5, 16), lambda i, j : C[i, j]*2, name='D')
s = tvm.create_schedule(D.op)
d_o, d_i = s[D].split(D.op.axis[1], factor=8)
s[C].compute_at(s[D], d_i)
C 的 attach path 是 {j_inner, j_outer, i}。IR 如下所示：
Building an Attach Path
按自顶向下的顺序遍历 D 的叶变量，从 C.attach_ivar 开始的，以及更低的叶变量都需要添加到 C 的 attach path 中。如果 D 也被 attach 到其他 stage E，那么也要对 E 的叶子进行同样的操作，直到没有其他 attach 的 stage。
C = tvm.compute((5, 16), lambda ci, cj : tvm.const(5, "int32"), name='C')
D = tvm.compute((5, 16), lambda di, dj : C[di, dj]*2, name='D')
E = tvm.compute((5, 16), lambda ei, ej : D[ei, ej]*4, name='E')
s = tvm.create_schedule(E.op)
s[C].compute_at(s[D], D.op.axis[1])
s[D].compute_at(s[E], E.op.axis[1])
D 的 attach path 是 {ej, ei}，而 C 的是 {dj, di, ej, ei}。IR 如下所示：
for ei 0, 5
for ej 0, 16
for di 0, 1
for dj 0, 1
for ci 0, 1
for cj 0, 1
C[0, ci + cj] = 5
D[0, di + dj] = C[0, di + dj] * 2
E[ei, ej] = D[0] * 4
InferBound with compute_at
InferRootBound 的目标是确定特定 stage C 的 root_iter_var 的 rang，其 Phase 1-2 给 C 的 consumer 的叶子 IterVar 设定 IntSet，然后将这些 IntSet 传播到 consumer 的 root_iter_var。如果不存在 attach， 为 consumer 的变量计算出的 range 定义了 consumer 需要多少 C。然而，如果这个 stage 实际上位于 consumer 的变量 j 的 scope 中，那么它一次只需要 j range 中的一个点。
Phase 1：初始化 consumer leaf_iter_vars 的 IntSet
有三种情境：
- Case 1：叶子变量的 range 为 1，这种情况下，等于 range 最小值
- Case 2：不需要 relaxation，只有一个数据点。当存在 compute_at 时，会出现该情况，将会对 attach_ivar，以及更高层的叶变量应用 Case 2，可保证只请求叶变量范围中的单个节点；
- Case 3：需要 relaxation，叶子的 range 被转换为 IntSet。
Phase 2：将 IntSet 从 consumer leaves 传播到 consumer 的 root
首先调用 PassUpDomain，直到 root 节点，然后创建 root_iter_vars 到 IntSet 的映射。在 PassUpDomain 的过程中， 会访问 consumer stage 的 IterVarRelation。
若是 split relation，需要根据 inner 和 outer 的 IntSet 设置父 InterVar 的 IntSet，具体而言：
- Case 1：如果 outer 和 inner 的 IntSet 和 range 是匹配的。那么父变量的 range 直接转换为 IntSet。
- Case 2：否则，使用
outer * rmap[inner]->extent + inner + rmap[parent]->min。
若 schedule 中包含 compute_at，就可能应用 Case 2。这是因为叶子的 IntSet 可能会被设置为单个点，因此 IntSet 将不再与 Range 匹配。如果 stage 并未被 attach 到当前 consumer，那么需要将对于这个 consumer 的 attach path 的每个变量 iv 的 range 添加到 relax_set。该 stage 的 root 变量需要根据这个 relax_set 进行计算。
例如，如下代码中，C 并未被 attach，但是 C 的 consumer D 在 E 中 attach 了。因此计算需要多少 C 时，必须考虑 D 的 attach path {ej, ei}。