PR Category

CINN

PR Types

Performance

Description

通过循环条件的蕴含（Entailment）关系对不整除的循环进行优化

这个PR实际上实现了两个优化，第一个优化是消除不整除循环里面的if，例如对于下面的循环：

for (k = 0; k < 8; k += 1) {
  if (((k * 32) + thread.x) < 234) {
    var_1[0] = var_1[0] + var[((block.x * 234) + (k * 32)) + thread.x];
  }
}

观察到，if的条件事实上蕴含（Entail）了循环的条件，即当((k * 32) + thread.x) < 234为真时，k < 8一定也为真。可以使用反证法证明，假设k < 8不为真，例如k = 8，则无论thread.x等于多少，((k * 32) + thread.x)都至少为256，不满足小于234的条件，因此蕴含关系成立。

因此，可以使用if的条件替代掉循环的条件，并把if消去，得到新的循环：

for (k = 0; ((k * 32) + thread.x) < 234; k += 1) {
  var_1[0] = var_1[0] + var[((block.x * 234) + (k * 32)) + thread.x];
}

完成第一个优化后，又注意到第二个优化，即此时k * 32是这整个循环中的公共表达式，可以提取出来变成k_strided = k * 32，这样进一步省掉一个寄存器，最终优化结果为：

for (k_stride = 0; (k_stride + thread.x) < 234; k_stride += 32) {
  var_1[0] = var_1[0] + var[((block.x * 234) + k_stride) + thread.x];
}

关于为什么两个优化实现在一个Pass里，因为这两个优化是息息相关的，都是和循环变量有关，两个优化都要对loop的extent和step进行修改，所以实现在一个Pass里更方便；另外优化二单独做是没有收益的，必须和优化一结合才有收益，因此分开实现可能导致优化一没匹配但优化二自己做了的情况

注：优化一是可以单独做的，在性能收益中占大头；优化二确实有可能匹配不上（尤其在动态shape下），但没关系，收益的大头已经拿到了

关于实现

由于新循环的条件是一个表达式而不是k < extent的形式，需要PolyFor才能表示，而当前我们的新IR中把PolyFor移除了，因此为了尽快上线该Pass，我只能先用一个C的宏来进行trick，即：

#define CINN_ENTAIL_LOOP_CONDITION(__loop_var, __cond, __stride) \
  }                                                              \
  for (decltype(__stride) __loop_var = 0; __cond; __loop_var += __stride) {

上述循环实际生成的IR为：

for (k = 0; k < 8; k += 1) {
  int k_strided = k * 32;
  CINN_ENTAIL_LOOP_CONDITION(k_strided, ((k * 32) + thread.x) < 234, 32);
  if ((k_strided + thread.x) < 234) {
    var_1[0] = var_1[0] + var[((block.x * 234) + (k * 32)) + thread.x];
  }
}

也就是用宏把原有循环打断，然后重新构造一个PolyFor的结构，这样确实不太规范，但至少在现行IR中是安全的，而且字面上没有改变循环的结构，调试的时候直接把CINN_ENTAIL_LOOP_CONDITION这行删掉也能正常执行

性能测试

• BatchNorm NHWC [125, 50, 50, 256]

• BatchNorm NCHW [-256, 256, -56, -56]（负号为动态Shape）