Halide Notes-算法和计算解耦

因为调研深度学习的端上加速方法, 了解到了 Halide, 顺便看了一下 Halide 的论文.

直接说 Halide 是计算加速的方案并不准确, Halide 更关注的是跨平台的计算加速, 包括 x86体系, arm 体系, cuda 体系等. Halide 的主要工作是把算法(algorithm)和具体的计算(schedule)做了解耦. 在这个场景下, 工程师做计算优化的过程是, 首先, 完成算法的定义, 然后, 对定义算法的 function 探索各种不同的 schedule. 因为 algorithm 和 schedule 是解耦的, 对于一个 algorithm, 不管使用哪种 schedule, 都不会影响 algorithm 的正确性, 因此, 工程师就可以非常大胆灵活地去尝试各种不同的 schedule. 另外, Halide 对于优化 pattern 做了一定的封装, 因此, 也大大减少了探索各种不同 schedule 的工作量. 最后, Halide 是一种 embedding language, 依附于 C++.

  • algorithm: 简单说就是我们的算法.
  • schedule: 算法在具体硬件平台上执行时的内容, 例如存储分配, 计算顺序等.
1
2
3
4
5
6
7
8
9
10
11
12
Func halide_blur(Func in) { Func tmp, blurred;
Var x, y, xi, yi;
//定义 algorithm
tmp(x, y) = (in(x-1, y) + in(x, y) + in(x+1, y))/3;
blurred(x, y) = (tmp(x, y-1) + tmp(x, y) + tmp(x, y+1))/3;
// 定义具体的 schedule, 根据平台的不同可以尝试各种不同的 schedule.
// tile, vectorize, parallel, chunk 是Halide定义的几种优化操作
blurred.tile(x, y, xi, yi, 256, 32)
.vectorize(xi, 8).parallel(y);
tmp.chunk(x).vectorize(x, 8);
return blurred;
}

计算优化中需要考虑重要一点是 locality, 在计算完 A 值之后, 如果我后面还要用到 A, 那么, 希望在用到 A 的时候可以在 cache 中找到 A, 而不是被换出去了.

Algorithm

Algorithm 定义了执行什么.
Halide 中, 图像是用 Function 来描述的, 例如上面例子中的 Func in, 函数在一点的值表示图像在该点的 color.
图像处理过程是多个 function 组成的 chain 过程.
Function 可以是简单的关于其 argument 的表达式, 也可以是有限域上的 reduction, 第二种情况的例子如计算直方图.

Schedule

Schedule 定义了怎样执行. 为了把计算做到极致, Schedule 有时候需要具备专业知识的人来精细调整. 但是, 不影响 Algorithm.
Halide 中有 4 中 schedule, Inline, Root, Chunk, Reuse

Inline: 在需要的时候计算, 结果不保存. Inline 存在的问题是, 假如在一个 pixel 的计算结果在多个位置用到, 那么就会做多次相同的计算, 浪费了计算济源. 但是, 从另外一个角度讲, 在图像处理的场景下, 速度有时候会被 bandwidth 限制住, 这个时候, 还不如直接使用寄存器中的数据再计算一遍来的快.

Root: 这个方法和 Inline 恰好是两个极端情形, Root 是在 caller 进行计算之前, callee 计算完所有的 caller 需要的数据. 例如在上面滤波的例子中, 先计算完所有的横向的数据, 然后再计算纵向的数据. 这个策略的好处很明显, 所有的点只会计算一次, 节省的计算开销, 但是, 问题是, 在计算domain比较大的情况下, 会破坏 storage 和 locality, 用来保存中间结果的 buffer 可能要求很大, 在最终caller 需要的时候, callee 中的 point 很可能已经不在 cache 中了.

Chunck: 处于上述两者之间, callee 先计算一个 subdomain 的结果, 然后 caller 利用这些结果, 然后 callee 计算下一个 subdomain 的结果, 然后 caller 继续计算. 这里需要做好 producer-consumer 之间的 trade-off.

Resue: 这种策略很简单, 就是 caller A 使用完 callee A 的结果之后, callee A 的结果并不删除, caller B 继续使用. 这里要求, callee A 要在 caller B 的 scope 中.

未完不待续……