[SCHEDULE] Improve bound inference for split

[merrymercy]

When we split an axis i with a factor f that is larger than the extent of i, we should set the extent of the inner axis to be the extent of i instead of f.

Test script

import tvm

A = tvm.placeholder((2, 2), name='A')
B = tvm.compute((2, 2), lambda i, j : A[i][j] + 1.0, name='B')

s = tvm.create_schedule([B.op])

i, j = s[B].op.axis
s[B].split(i, 4)

print(tvm.lower(s, [A ,B], simple_mode=True))

Results

Before

produce B {
  for (i.inner, 0, 4) {
    for (j, 0, 2) {
      if (likely((i.inner < 2))) {
        if (likely((i.inner < 2))) {
          B[((i.inner*2) + j)] = (A[((i.inner*2) + j)] + 1f)
        }
      }
    }
  }
}

After

produce B {
  for (i.inner, 0, 2) {
    for (j, 0, 2) {
      B[((i.inner*2) + j)] = (A[((i.inner*2) + j)] + 1f)
    }
  }
}

[merrymercy]

cc @jcf94 @tqchen

[yongfeng-nv]

@merrymercy, I submitted almost the same PR last week: #4885. More failures are exposed during CI. I am improving it with two failures left: #4931. Can you review it?

[merrymercy]

@yongfeng-nv
Good work! I will review your PR and close this one

[yongfeng-nv]

@merrymercy, thanks. We ran into this issue, because the 'if' statements are not allowed in the tensorization region, giving errors like: TVMError: Tensorize failed, split condition likely((i.inner < 8)) relies on var defined inside tensorize scope.

Can you share your use case stopped by this issue? I am going to open a thread in the forum about it and finding more use cases while fixing failures.