Open jjsjann123 opened 1 month ago
Could this be solved by accepting slice-only segments (and metadata-only reshapes) in the ExprEval
scheduler?
Could this be solved by accepting slice-only segments (and metadata-only reshapes) in the
ExprEval
scheduler?
That's a great alternative.
My wild guess is that ExprEval
scheduler is not populating/updating output TensorView definition to its consumer segments, so there's still that part of plumbing work to be done?
(slice from ExprEval
might update the output tensor to be not contiguous... unlike the default behavior where everything is contiguous...)
My wild guess is that
ExprEval
scheduler is not populating/updating output TensorView definition to its consumer segments, so there's still that part of plumbing work to be done? (slice fromExprEval
might update the output tensor to be not contiguous... unlike the default behavior where everything is contiguous...)
Oh that's a good point; this is the real challenge in mixing the alias analysis with segmentation isn't it. We could envision a system where during segmentation, the a scheduler can optionally set some options on its output tensors when it accepts a segment, so ExprEval could set the contiguity flags at this point. If the allocation domain is already set on those outputs, this call could fail and the scheduler could either adjust or reject the segment, resetting the output TV.
I believe below is what's happening at this moment:
alias_to_source_
, which is good. https://github.com/NVIDIA/Fuser/blob/1026f676b370a7ca5c6472daf1de061c8755bc87/csrc/alias_analysis.cpp#L273Once we fix (2), NoOpScheduler should be able to pick up the single-slice segment. Caveat: it may be risky to aggressively set the allocation domain of an intermediate tensor; I don't know how schedulers treat them.
where during segmentation, the a scheduler can optionally set some options on its output tensors when it accepts a segment,
would be great though it's a bigger hammer.
Caveat: it may be risky to aggressively set the allocation domain of an intermediate tensor; I don't know how schedulers treat them.
Yeah this part is also a bit concerning to me.... It felt like it's hard to make this decision on an unsegmented fusion.
In full rope example, we have a fusion starting with something like this:
nvfuser unfortunately ended up segmenting on
slice
, which gives us a segment like this one below, where we are missing out on leveraging alias analysis but generated a kernel for a single slice op.The issue here is that, we might not always segment on the right position. Current alias analysis runs as preseg pass, so it will miss out alias on segments and we are unfortunately generating kernels for the example above.
repro code
Note this needs to run with bookend disabled in thunder, as in @wujingyue 's PR: https://github.com/Lightning-AI/lightning-thunder/pull/731