General broadcast addition for source (C) tensor

[masahi]

Hi, I have a matmul config where the result of A * B is of shape (2, 64, 64, 320) and the C tensor is of shape (2, 1, 1, 320). I want to do A * B + C with such shapes, where the addition has the broadcasting semantics. Such input appears in Stable Diffusion. This pattern appears as part of a residual block, so I want to be able to do this using DefaultConv2dFpropWithBroadcast + LinearCombinationResidualBlock.

Does CUTLASS support such general broadcast addition for the source tensor? I quickly tried but it fails at runtime or yields garbage output.

@hwu36

[hwu36]

what are batch, m, n, k in your shapes?

[masahi]

yeah sorry, this is conv2d not matmul. So the data shape is (2, 64, 64, 320) and the filter is (320, 3, 3, 320).

If the source tensor is a bias of shape (1, 1, 1, 320), I can use the "stride 0" trick. The question is what if the source tensor is batched. The addition in A * B + C is no longer a simple elemwise operation anymore, but such broadcasting addition is widely supported by pytorch, numpy etc.

[hwu36]

so different batch wants to load different vector?

[masahi]

Yes conceptually, but rather than adding a special case logic to handle such case, I wonder if it is a good idea to support general broadcasting math like the rest of numerical computing framework like numpy / pytorch / tvm etc already do.

[hwu36]

instead of setting all strides to 0, you could set it them as (320, 0, 0). So, nothing happens when H or W increase, but when N increases, you will move to the next batch.