[SCHEDULE] Generate Lower Bound Conditions

[derisavi-huawei]

given the following algorithm and schedule

import tvm
A = tvm.placeholder((64,), name='A', dtype="float32")
Apad = tvm.compute((66,), lambda i: tvm.select(tvm.all(i>0, i < 65), A[i-1], tvm.const(0.)), name='Apad')
Apad2 = tvm.compute((66,), lambda i: Apad[i]*2, name='Apad2')
s = tvm.create_schedule(Apad2.op)
AL1 = s.cache_read(A,"local",[Apad])
print(tvm.lower(s, [A, Apad2], simple_mode=True))

I get the following lowered code. The stage that produces A.local seems incorrect since it accesses element A[-1] in the first iteration. Is there a solution or a workaround?

// attr [A.local] storage_scope = "local"
allocate A.local[float32 * 66]
// attr [Apad] storage_scope = "global"
allocate Apad[float32 * 66]
produce A.local {
  for (ax0, 0, 66) {
    if (likely((ax0 < 65))) {
      A.local[ax0] = A[(ax0 + -1)]
    }
  }
}
produce Apad {
  for (i, 0, 66) {
    Apad[i] = tvm_if_then_else(((0 < i) && (i < 65)), A.local[i], 0.000000f)
  }
}
produce Apad2 {
  for (i, 0, 66) {
    Apad2[i] = (Apad[i]*2.000000f)
  }
}

[tqchen]

this is a problem, we should have add guard for the beginning condition of A as well

[kun-zh]

Seems the size of A.local should be 64, and don't need to generate the condition assignment. Am I right?

[kun-zh]

Sorry, the size for A.local is 66 because TVM will schedule A.local with the scheduling for the reader tensor. However, we want to get the expected lower code:
// attr [A.local] storage_scope = "local"
allocate A.local[float32 * 66]
// attr [Apad] storage_scope = "global"
allocate Apad[float32 * 66]
produce A.local {
for (ax0, 0, 64) {
A.local[ax0] = A[(ax0)]
}
}
produce Apad {
for (i, 0, 66) {
Apad[i] = tvm_if_then_else(((0 < i) && (i < 65)), A.local[i + -1], 0.000000f)
}
}
produce Apad2 {
for (i, 0, 66) {
Apad2[i] = (Apad[i]*2.000000f)
}
}

I tried:
A = tvm.placeholder((64,), name='A', dtype="float32")
AL = tvm.compute((64,), lambda i: A[i], name='AL')
Apad = tvm.compute((66,), lambda i: tvm.select(tvm.all(i>0, i < 65), A[i-1], tvm.const(0.)), name='Apad')
Apad2 = tvm.compute((66,), lambda i: Apad[i]*2, name='Apad2')
s = tvm.create_schedule(Apad2.op)
s[AL].set_scope("local")
#AL = s.cache_read(A,"local",[Apad])
print(tvm.lower(s, [A, Apad2], simple_mode=True))

And get the error:
raise ValueError("Cannot find the operation %s in schedule" % (str(k)))
ValueError: Cannot find the operation compute(AL, 0x15ff280) in schedule

Seems AL has been removed because of the memory optimization. I think if the scopes are different between A and AL, AL should not be removed.

[derisavi-huawei]

@kun-zh
The reason for the error message is that AL is not used any of the computations. If you replace A in Apad's expression with AL in your code, then your code compiles but it will still have a similar problem. The lowered code is

// attr [AL] storage_scope = "local"
allocate AL[float32 * 66]
// attr [Apad] storage_scope = "global"
allocate Apad[float32 * 66]
produce AL {
  for (i, 0, 66) {
    if (likely((i < 65))) {
      AL[i] = A[(i + -1)]
    }
  }
}
produce Apad {
  for (i, 0, 66) {
    Apad[i] = tvm_if_then_else(((0 < i) && (i < 65)), AL[i], 0.000000f)
  }
}
produce Apad2 {
  for (i, 0, 66) {
    Apad2[i] = (Apad[i]*2.000000f)
  }
}