When used in conjunction with tensors that have a dimension of size 1 topi.concatenate will result in lowered code that has a negative (invalid) iteration bound.
Example:
import tvm
from tvm import relay
from tvm import te
import topi
shape = (4, 2, 3)
shape_one_layer = (4, 3)
data = te.placeholder(shape)
l0 = te.compute(
shape_one_layer,
lambda x, y: data[x, 0, y] + 1,
name='layer0')
l1 = te.compute(
shape_one_layer,
lambda x, y: data[x, 1, y] + 1,
name='layer1')
l0e = topi.expand_dims(l0, 1, 1) # new shape (4, 1, 3)
l1e = topi.expand_dims(l1, 1, 1) # new shape (4, 1, 3)
result = topi.concatenate((l0e, l1e), axis=1) # shape (4, 2, 3)
s = te.create_schedule(result.op)
print(tvm.lower(s, [data]))
produces:
allocate(T_concat, float32, [24]) {
for (x: int32, 0, 4) {
for (y: int32, 0, 3) {
layer1[((x*3) + y)] = ((float32*)placeholder_2[(((x*6) + y) + 3)]) + 1f32)
}
}
for (x_1: int32, 0, 4) {
for (y_1: int32, 0, 3) {
layer0[((x_1*3) + y_1)] = ((float32*)placeholder_2[((x_1*6) + y_1)]) + 1f32)
}
}
for (ax0: int32, 0, 4) {
for (ax1: int32, 0, 2) {
for (ax2: int32, 0, 3) {
T_concat[(((ax0*6) + (ax1*3)) + ax2)] = @tvm_if_then_else((1 <= ax1), (float32*)layer1[((((ax0*3) + (ax1*3)) + ax2) - 3)]), (float32*)layer0[(((ax0*3) + (ax1*3)) + ax2)]), dtype=float32, type="pure_intrin")
}
}
}
}
}
The negative bound occurs on the line where T_concat
is computed, in the way layer1
is accessed: layer1[((((ax0*3) + (ax1*3)) + ax2) - 3)]
. When starting iteration all three variables are 0, resulting in an index of -3
.
A workaround is to use if-then-else directly:
result = te.compute(
shape,
lambda x, t, y: te.if_then_else(t == 0, l0e[x, 0, y], l1e[x, 0, y])) # shape (4, 2, 3)
s = te.create_schedule(result.op)
print(tvm.lower(s, [data]))
produces the correct iteration bounds:
compute[(((x_2*6) + (t*3)) + y_2)] = @tvm_if_then_else((t == 0), (float32*)layer0[((x_2*3) + y_2)]), (float32*)layer1[((x_2*3) + y_2)]), dtype=float32, type="pure_intrin")
I don’t know the source of this problem. Opinions are welcome. Also, is there any better way of concatenating tensors in TVM?