expand concretization to support opaque handling of stride order and shape specialization

[jjsjann123]

Proposal

Add a concretization-like step to mutate TensorView definition with correct stride_order (allocationDomain) and contiguity flag.

Context

currently nvfuser API requires user to statically define broadcast dimensions as well as stride order (not even in our integration yet!) at the time when a fusion is defined. Any further update to the tensor shape / stride that contradicts these definition requires user/integration to handle that with a custom cache strategy.

It is challenging to get a custom cache right, because of the complexity involved. The caching strategy in nvfuser evolves over time, which makes it even harder to maintain.

We should try to keep this complexity inside nvfuser and abstract away such definition on user API.

i.e. when we define a TensorView, we can specify only the rank of a tensor, without specifying the stride order or broadcast semantics. At runtime when a concrete input tensor is given, we can concretize the definition of each tensor, propagate that information to have a concrete fusion that can be scheduled.

Eventually we would want to use this to handle both stirde_order as well as shape specialization (broadcast / size-0 or whatever new thing that pops up).

TODO

1. We can take the initial step to support stride_order, since it's simpler to add and doesn't require propagation. It works very much like the reshape handling in concretization and is somewhat isolate from the rest of the cache system we have in our stack.
2. The protocol for how to interpret a stride-order can be tricky and that propagation (if we want to do) is even harder. (for some fun context, look at these: Memory Format support for Resnet models pytorch/pytorch#23403 suggest_memory_format has ambiguity & cannot represent intended layout format for corner cases pytorch/pytorch#24090) But luckily once those are inside codegen logic, changing them would be relatively easy. So we don't have to make the right decision in the first place.

[jjsjann123]

cc'ing @wujingyue

[jjsjann123]

Also, wondering what @naoyam @zasdfgbnm @jacobhinkle thinks about this one.

[jacobhinkle]

Could this happen at the FusionCache level? That is, we could have fd.define_tensor(ndims=4, dtype=DataType.Float) where all we are able to say is the dtype and dimension. Then execute would take the inputs and infer how to translate the corresponding TensorRecords to TensorViews. I guess we'd also need to take inputs into account during hashing/comparison which could impact latency but should be better than handling it in integration...

Alternatively, putting it into FusionExecutorCache means we'd need to be able to express ambiguity in stride and IterType, and then be able to properly propagate them, which seems possible but complex.

[wujingyue]

Thanks for writing this up, @jjsjann123 ! After talking to @kevinstephano , I'll be tackling the segmenter for the short term and unfortunately won't have time immediately for this. But it does sound like a fun and important problem to solve and I hope I can dive into this much more in the future.

[jjsjann123]

Could this happen at the FusionCache level?
Alternatively, putting it into FusionExecutorCache...

I'm a bit unsure about FusionCache, IIUC, that's going to be available to only python API?
putting it inside FusionExecutorCache is universal to both python and cxx api, which makes me slightly happier about it~~

[jacobhinkle]

Yeah, I think it's doable to build into FusionExecutorCache.

Handling unknown broadcasts is pretty simple: we would just mark the IterTypes as Symbolic in TensorRecord and concretize inputs the same way we do for resized IterDomains currently.

Concretizing contiguity should only apply to input tensors so we could just reset contiguity of inputs at concretization.

Each of these two steps looks pretty straightforward actually.

[kevinstephano]

@jacobhinkle or @jjsjann123 is this still an issue?

[jacobhinkle]

We have not handled it yet. I don't think it's extremely hard to do on the nvFuser side. Maybe a 2-3 week effort or month long task for a new contributor. @jjsjann123 do you think Thunder would benefit from nvFuser handling this level of caching?

[jjsjann123]

I think this is a nice thing to have for our old API. But as we move towards direct binding, I'm not sure giving user more flexibility when defining fusion would be beneficial. This feels like adding more complexity.

I'm closing as not-planned.