[FeatureRequest] Pad on broadcast dimensions are not well supported by codegen (both frontend / backend)

[jjsjann123]

Background

This is a pretty involved topic. The original question rises from @jacobhinkle 's comment on the lack of support on padding broadcast dimensions from codegen. #10 (comment)
I had an offline discussion with @naoyam on this and want to open up the issue to track the problem.

1. We need to support padding on broadcast dimensions. Look at the example below vvv:

def fn0():
  t0 = torch.randn(1)  # how are we define t0 IterDomain?                       
  t1 = torch.randn(5)
  o1 = torch._C._nn.pad(t0, (2, 1))
  o2 = t0 + t1

t0's IterDomain needs to be marked as broadcast, since we need to resolve output shape of o2 from t0 + t1, where t1 has a non-broadcast IterDomain.
Meanwhile, t0 is being padded, so we can't really work around it without proper backend support.

2. Our current API seems to suggest that we support dynamic pad length.
   

   Fuser/csrc/ops/alias.cpp

   Lines 382 to 384 in 86d5dd3

   	// Padding widths are assumed to be non-negative. Currently there's no
   	// validation.
   	TensorView* pad(TensorView* inp, const std::vector<Val*>& pad_widths) {

I doubt that we actually can support it.

def fn1(pad_left, pad_right):
  t0 = torch.empty([1])
  t1 = torch._C._nn.pad(t0, (pad_left, pad_right))  # how are we define t1 IterDomain?

Given the example above, t0 has size-1 rank-1, which means it comes with a broadcast IterDomain.

So the tricky part here comes to, what do we do with t1? Does t1 has an broadcast or non-broadcast IterDomain? Well, that depends on pad_left/pad_right value!
Imagine if you are given 0 pad at runtime, then the output t1 needs to have a broadcast IterDomain, while with non-zero padding, we'll have a non-broadcast IterDomain (Note: no it can't be a broadcast with extend, since we are padding a certain value here!).

3. The last point is trivial. Currently codegen pad assumes non-negative padding value. Which is unfortunately not the case for pytorch.

>>> x = torch.randn(5)
>>> o = torch._C._nn.pad(x, (-2, -2))
>>> o
tensor([-0.6133])

Proposal

A few things we think that is needed:

1. We should support pad on broadcast dimensions.
2. We should enforce non-negative pad length on the op and have frontend transform pad with negative length to slice or other ops to keep our lives simple.
3. Maybe we should limit pad length to be static, that would make it straightforward to plumb the support of pad on broadcast dimensions. (In combination with the non-negative pad length, we'll know at compile time whether the output IterDomain should be broadcast or not.)

Note that point 3 here is not a decision that we can make here without looking at the models that we want to support. cc'ing @kevinstephano Do you think for nanogpt, static size padding is sufficient for now?

Tracking things:

• is static pad length sufficient for our initial support?
• does the same problem translate to cat / slice?

[jacobhinkle]

It seems like this really calls for inferring a property on that IterDomain based on later uses within the Fusion which is a topic that comes up for dynamic shapes with other ops like view also. For instance we might mark the ID as "MaybeBroadcast" then when it is used later we would infer some potentially complicated relations between the inputs and the sizes and types of IDs. For pointwise binary ops the matching domains would imply that they match or at least one is a broadcast, for example. When we have processed the whole graph we can try and satisfy those relations and proceed as usual. So this issue is related to IterDomain graphs but with the new wrinkle that we may need to change the graph by changing whether to mark as bcast.

[naoyam]

It seems like this really calls for inferring a property on that IterDomain based on later uses within the Fusion which is a topic that comes up for dynamic shapes with other ops like view also. For instance we might mark the ID as "MaybeBroadcast" then when it is used later we would infer some potentially complicated relations between the inputs and the sizes and types of IDs. For pointwise binary ops the matching domains would imply that they match or at least one is a broadcast, for example. When we have processed the whole graph we can try and satisfy those relations and proceed as usual. So this issue is related to IterDomain graphs but with the new wrinkle that we may need to change the graph by changing whether to mark as bcast.

Exactly. Here's what I'm prototyping: https://github.com/NVIDIA/Fuser/compare/dynamic_reshape. I added a new IterType called Symbolic.

[jacobhinkle]

This functionality has existed for some time now.