We are trying to use TVM to generate operator definitions for MXNet.
The gap is, despite the fact that TVM compute/schedule can deal with symbolic shape, some compute definition strongly rely on a fixed shape. For example, broadcast ops,
A = tvm.placeholder(shape=(tvm.var("a1"), tvm.var("a2")))
B = tvm.placeholder(shape=(tvm.var("b1"), tvm.var("b2")))
C = A + B
Currently topi’s broadcast implement will not work as it needs to check whether each dim equal to ONE. e.g., https://github.com/dmlc/tvm/blob/master/topi/include/topi/detail/broadcast.h#L59
While in the case of AOT for MXNet (instead of JIT), we don’t know the shape at the compile time.
My proposal is, introduce a new type of buffer (e.g., broadcast buffer), insert stride[i] = (shape[i] == 1) ? 0 : stride[i] in arg_binder.cc, for example,
C[i * s1 + j] = A[i * sa1 + j] + B[i * sb1 + j]
if shape(A) = [10, 3] and shape(B) = [1, 3], then the previous code will set sb1 = 0, which will be correct.
And the compute definition can be,
m0, m1 = tvm.var("m0"), tvm.var("m1")
n0, n1 = tvm.var("n0"), tvm.var("n1")
o0, o1 = tvm.var("o0"), tvm.var("o1")
A = tvm.placeholder((m0, m1), name='A', dtype=?)
B = tvm.placeholder((n0, n1), name='B', dtype=?)
C = tvm.compute((o0, o1), lambda i, j: A[i, j] + B[i, j], name='C')
Ab = tvm.decl_buffer(A.shape, A.dtype, type="broadcast", strides=[tvm.var("sa"), 1])
Bb = tvm.decl_buffer(B.shape, B.dtype, type="broadcast", strides=[tvm.var("sb"), 1])
s = tvm.create_schedule(C.op)
f = tvm.lower(sch, [A, B, C], binds={A:Ab, B:Bb})