Allow pipelining loads that don't feed directly into tl.dot.

[jlebar]

Allow pipelining more loads indirectly used by a dot.

Previously if a load was used indirectly by a dot but there were non-unary
operations between the load and the dot, we couldn't pipeline it.

Also clean up some TODOs.

PR chain

1. 👉 Allow pipelining loads that don't feed directly into tl.dot. #3415 👈 YOU ARE HERE

[ThomasRaoux]

I didn't look at all the changes but there is already annotations that can be set on the loop to force pipelining. I think this is what we would want to ise for this case instead of adding to the existing logic that detects load + dot pattern.
Here is the PR that enabled it:
#3092

[jlebar]

We are already passing num_stages as an argument to the kernel. It appears this does the same thing as passing num_stages to the for loop's tl.range op. https://github.com/openai/triton/blob/8fed348280c32ef22c0e44b69bcd145420f78710/python/triton/compiler/code_generator.py#L933

To be clear, the issue isn't that the loop is not pipelined, it's that this particular load within the loop is not pipelined, because it does not feed into the dot directly.

Am I misunderstanding something?

[ThomasRaoux]

We are already passing num_stages as an argument to the kernel. It appears this does the same thing as passing num_stages to the for loop's tl.range op.

https://github.com/openai/triton/blob/8fed348280c32ef22c0e44b69bcd145420f78710/python/triton/compiler/code_generator.py#L933

To be clear, the issue isn't that the loop is not pipelined, it's that this particular load within the loop is not pipelined, because it does not feed into the dot directly.

Am I misunderstanding something?

In addition to passing a per-loop num_stages the idea was to also bypass the heuristic that only pipelines matmul loop. So I would think we would pipeline every load that can be pipelined in this case. It's possible the PR I pointed out behaves differently than what I thought. I do think it would be nice to not have to rely on load feeding into a dot for this case. (of course unless we mark explicitly ops there will always be a part of heuristic) but I though the "general" case would do what you want here.
I could be wrong and the code may be missing something. If possible I think we should avoid changing the heuristic.

For context the heuristic is there for both historic and user simplicity reasons. Historic because we use to only have the global control and the compiler had to pick the loop to pipeline and for simplicity as the existing pattern checks usually always benefits from being pipelined.

[jlebar]

Updated the PR per our discussion

[jlebar]

Thank you for the review. :)