[RFC][Codegen] Broadcast ops with symbolic shape

[yzhliu]

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')
B = tvm.placeholder((n0, n1), name='B')

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: #3389

@tqchen @reminisce @junrushao1994

[yzhliu]

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.

[junrushao]

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

[tqchen]

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

[junrushao]

@were

[yzhliu]

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?

[junrushao]

I am in favor of migrating as well :-)

Ping @were again for another round of review.

[tqchen]

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

[yzhliu]

close by #3389