Closed KuangjuX closed 12 months ago
由于观测到内存操作在模型推理运行中占据了大量的时间,算子的输入输出在三级内存(寄存器,共享内存,全局内存)中互相搬运浪费太多时间,因此想提出一个方案用于解决过多的内存加载与存储。
Welder 提出了 tile-graph 的概念,采取先连接后融合。首先猜测两个相邻的算子是否可以在某个内存层级进行数据 tile 重用,然后获得最通用的 tile shape 查看是否可以减少内存流量。抽象了 SetConnect, Propagate 两个接口。设计了两步算法,图连接和子图调度。
Conv 算子计算 1x1xC 大小的 tile,输入是一个 3x3xC 的 data tile,MaxPool 输入一个 2x2xF,输出 1x1xF。为了提高内存利用资源,Welder 允许两个相邻的算子 tile 通过中间 tile 进行连接,也叫做 reuse-tile。
在 Figure 5 中 Conv 和 ReLu 在 L0 进行融合并形成了一个 tile-graph。同时它们在 L1 级别被视为一个虚拟节点。ReLU 的输出被 spill 到 L1 中并且作为 MaxPool 的输入。这允许三个算子在 L1 形成一个单 tile-graph。
为了实际在硬件加速器生成代码,Welder 提出了几个内存相关的抽象。Figure 8 展示了一个执行多层次 tile-graph 使用这些接口的流程。
TVM: compute_inline
调度原语用来实现寄存器级别 tile 连接。对于 sharded memory 级别仅仅生成单个 kernel 并使用几个额外的 passes 去将这些单独的 kernels 压缩成一个融合内核。
略。
paper: https://www.usenix.org/system/files/osdi23-shi.pdf