NVIDIA / Fuser

A Fusion Code Generator for NVIDIA GPUs (commonly known as "nvFuser")
Other
228 stars 42 forks source link

allocation domain propagation to support reshape operations #2235

Open jjsjann123 opened 2 weeks ago

jjsjann123 commented 2 weeks ago

allocation domain propagation should support reshape operations. @zasdfgbnm brought it up here: https://github.com/NVIDIA/Fuser/pull/2168#discussion_r1592894415

Currently the propagation only maps exact iterdomains from sources to destinations. Handling reshape requires us to walk through the ValGraph and establish rules to map through split/merge/resize. Which might not be too bad to resolve if we are looking at the relationship between pair of source to destination. As @wujingyue suggested in the thread. @zasdfgbnm and @jacobhinkle were suggesting to establish a uniformed relationship table per fusion by traversing the IR. Which could used later by schedulers/segmenters to coherently establish allocation domain for boundary tensors.

I think the uniformed allocation domain sounds interesting but it's going to be hard(?!) to justify the approach at current stage. (how useful would it really be at solving problems at hand?) Meanwhile the source to destination mapping seems like a low hanging fruit that I think can be relatively easy/quick to pull of. I'm not ruling out the alternative yet. Opening the issue here to track the progress as well as bookkeeping the alternative that we are not picking up at the moment.