[TVM][BUGFIX] Fix missing reduction init predicates

[sgrechanik-h]

The commit 6fdc3e0 triggered the following bug. Consider this code:

    n = 10
    X = tvm.placeholder(shape=(10,), dtype='int32', name="X")
    k = tvm.reduce_axis((0, n))
    Y = tvm.compute((), lambda: tvm.sum(X[k], k), name="Y")

    s = tvm.create_schedule([Y.op])
    ax = s[Y.op].fuse(*Y.op.axis)
    axo, axi = s[Y.op].split(ax, nparts=20)

    print(tvm.lower(s, [Y, X], simple_mode=True))

Before 6fdc3e0 the generated code looked like this:

produce Y {
  for (singleton.outer, 0, 20) {
    Y[0] = 0
    for (rv, 0, 10) {
      Y[0] = (Y[0] + X[rv])
    }
  }
}

After 6fdc3e0 the code looks like this:

produce Y {
  for (singleton.outer, 0, 20) {
    Y[0] = 0
    for (rv, 0, 10) {
      if (likely((singleton.outer < 1))) {
        Y[0] = (Y[0] + X[rv])
      }
    }
  }
}

which is wrong, the result gets reset to zero. The reason is that the initialization is not guarded by the condition singleton.outer < 1.

The fix I propose is to skip loops related to the reduction axes instead of skipping loops unrelated to the compute op axes when generating predicates for the initialization code. The result now looks like this:

produce Y {
  for (singleton.outer, 0, 20) {
    if (likely((singleton.outer < 1))) {
      Y[0] = 0
    }
    for (rv, 0, 10) {
      if (likely((singleton.outer < 1))) {
        Y[0] = (Y[0] + X[rv])
      }
    }
  }
}

@tqchen @derisavi Please take a look.

[ZihengJiang]

why not flag == 2?

[sgrechanik-h]

Good point, this actually should be flag == 2. Added a test case for this.