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,
```python
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,
```python
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,
```python
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")
Bb = tvm.decl_buffer(B.shape, B.dtype, type="broadcast")
s = tvm.create_schedule(C.op)
f = tvm.lower(sch, [A, B, C], binds={A:Ab, B:Bb})
```
Related PR: https://github.com/dmlc/tvm/pull/3389
@tqchen @reminisce @junrushao1994
--
You are receiving this because you are subscribed to this thread.
Reply to this email directly or view it on GitHub:
https://github.com/dmlc/tvm/issues/3390
