Partition fix with rfactor, simplify and likely predicates.

[csarofeen]

Currently loop partitioning does not work with rfactor. This is simply because the predicate generated from the loop nest in rfactor is not wrapped in a likely so loop partition doesn't process it.

Currently

import tvm
K = 16*4+4

k = tvm.reduce_axis((0, K), 'k')

A = tvm.placeholder((1, K), name='A')

B = tvm.compute(
    (1,),
    lambda b:
        tvm.sum(A[b, k],
        axis=k),
                name='B'
)

s = tvm.create_schedule(B.op)
deps = [A, B]
ko, ki = s[B].split(s[B].op.reduce_axis[0], 16)
BF = s.rfactor(B, ko, 0)

s.normalize()
bounds = tvm.schedule.InferBound(s)
stmt = tvm.schedule.ScheduleOps(s, bounds)

stmt = tvm.ir_pass.LoopPartition(stmt, True)
stmt = tvm.ir_pass.Simplify(stmt)
print(stmt)

generates:

// attr [compute(B.rf, 0x5582e112a6c0)] realize_scope = ""
realize B.rf([0, 5], [0, 1]) {
  produce B.rf {
    for (k.outer, 0, 5) {
      B.rf(k.outer, 0) =0.000000f
      for (k.inner, 0, 16) {
        if ((k.inner < (68 - (k.outer*16)))) {
          B.rf(k.outer, 0) =(B.rf(k.outer, 0) + A(0, (k.inner + (k.outer*16))))
        }
      }
    }
  }
  // attr [compute(B, 0x5582e1128fe0)] realize_scope = ""
  realize B([0, 1]) {
    produce B {
      B(0) =0.000000f
      for (k.outer.v, 0, 5) {
        B(0) =(B(0) + B.rf(k.outer.v, 0))
      }
    }
  }
}

With this patch it will generate:

// attr [compute(B.rf, 0x559cdeecf6c0)] realize_scope = ""
realize B.rf([0, 5], [0, 1]) {
  produce B.rf {
    for (k.outer, 0, 4) {
      B.rf(k.outer, 0) =0.000000f
      for (k.inner, 0, 16) {
        B.rf(k.outer, 0) =(B.rf(k.outer, 0) + A(0, (k.inner + (k.outer*16))))
      }
    }
    B.rf(4, 0) =0.000000f
    for (k.inner, 0, 4) {
      B.rf(4, 0) =(B.rf(4, 0) + A(0, (k.inner + 64)))
    }
  }
  // attr [compute(B, 0x559cdeecdfe0)] realize_scope = ""
  realize B([0, 1]) {
    produce B {
      B(0) =0.000000f
      for (k.outer.v, 0, 5) {
        B(0) =(B(0) + B.rf(k.outer.v, 0))
      }
    }
  }
}

[csarofeen]

This seems to work well with a single rfactor, however, nested rfactors seem to still not be partion-able. Will keep working but I believe this PR is pretty low risk for unwanted side effects.

[yzhliu]

high-level question, IMO loop partition should work on any predicate, not only those with likely, right?