[dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Yizhi Liu <no...@github.com>]

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

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Yizhi Liu <no...@github.com>]

close by https://github.com/dmlc/tvm/pull/3389

-- 
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#issuecomment-507793738

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Tianqi Chen <no...@github.com>]

tag other folks who might have some thoughts/suggestions about this @ajtulloch @hlu1 @ZihengJiang @masahi @kazum 

-- 
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#issuecomment-503644825

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Tianqi Chen <no...@github.com>]

I am OK with the change. 
- Let us use enum as it will make things more clear (we could use str in python API)
- Let us make sure code simplifies in the auto-broadcast mode

Re: migration

We need to think about the implication of simplifiers. Given that auto_broadcast bind at a very late stage, the simplification will not be as powerful.  Need to make sure that the new approach does not compromise the original code when the shape is know




-- 
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#issuecomment-506786508

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Junru Shao <no...@github.com>]

I don't have preference between enum and string either.

-- 
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#issuecomment-503599102

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Yizhi Liu <no...@github.com>]

Regarding how it would affect topi implementation of broadcast and reduce ops, I think we can migrate current compute implement to this new approach, while we might still need to know the shape beforehand in schedule.

I'll change the name to "auto_broadcast" and document the behavior more precisely.

kindly remind everyone here to see if we can proceed?

-- 
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#issuecomment-506559414

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Junru Shao <no...@github.com>]

I am in favor of migrating as well :-)

Ping @were again for another round of review.

-- 
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#issuecomment-506610158

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Yizhi Liu <no...@github.com>]

Closed #3390.

-- 
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#event-2455283439

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Yizhi Liu <no...@github.com>]

Move the discussion here

> The things to be improved are
* Document the behavior independent of arg_binder.
  - Maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 0
* Debate on the name (auto broadcast?), enum vs string as type key
* Discuss how would the behavior affect topi implementation of broadcast and reduce ops
* Think about alternative implementations
  - Do we have to introduce stride?
  - e.g. j -> min(j, shape[ndim])
  - j -> j * scale(scale = 0 or 1)

Junru's comment,
> I do think it is proper to introduce stride. First, as far as we could see, using strides does not seem to have too much overhead, compared with other workarounds. Second, it seems to pave ways for other possible hacks to the buffer.

`j -> min(j, shape[ndim])` seems to introduce overhead, IMO `j -> j * scale(scale = 0 or 1)` is equally good, while we need to think where to allow user specify the behavior, I feel implicitly inserting such scale will cause surprise.
I don't have preference between enum and string.


-- 
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#issuecomment-503286893

Re: [dmlc/tvm] [RFC][Codegen] Broadcast ops with symbolic shape (#3390)

[Junru Shao <no...@github.com>]

@were 

-- 
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#issuecomment-503665088