Fix handling of unvisited operands in AxisInfoAnalysis

[neildhar]

We currently force initialisation of operands that have not yet been
visited with setToEntryState. This means that the order in which
values are visited can change the results of the analysis.

This can be a source of bugs. For example, the lowering for
AsyncCopyGlobalToLocalOp validates that the load addresses permit
sufficient vectorisation, however, this is up to the analysis actually
recovering the same information it had when the async copy was created.
Otherwise, we crash during lowering. I have an actual repro for this
but it has been very difficult to minimise it enough to make it suitable
for an lit test:
https://gist.github.com/neildhar/7eea6a312afa39d1cc83dc12627c2ba3

Populating the operands in this way also means that we have to handle
control flow like ForOp and IfOp explicitly in setToEntryState,
because we may attempt to populate their results when we visit their
users.

Instead, when we encounter an operation whose operands have not yet been
encountered, skip over the operation entirely. We can revisit it once
the operands have actually been visited. This improves the quality of
the analysis, and leaves the handling of control flow to the dataflow
framework.

[ThomasRaoux]

Why is it expected to be better? Is it possible to write a test where the results are better?
@Mogball, does that match what you expect?

[neildhar]

@ThomasRaoux Yes, although it has been very difficult to minimise into something that would be suitable for a test. The bug in the PR description is an example of this. Existing tests should already exercise the changes in this PR though, since we cannot delete the control flow handling without the fix to skip over operations with unvisited operands.

I think while we're operating on structured control flow, this probably is not any better. The framework visits all the operations in order, with the exception of the scf ops, but we special case them by initialising them with the best case. So forcing all the operands to be initialised should only affect scf ops, which we handle.

However, once the structured control flow is lowered out, the order in which things are visited becomes important. If we visit an operation before its operands are visited. We may end up pessimising the analysis, because we assign that operand the most conservative state. This is what happens in the bug in the PR description. Performing the analysis on the structured control flow in an earlier pass determines that we can use AsyncCopyGlobalToLocal. However, when we actually try to lower the async copy in LoadStoreOpToLLVM, we can no longer recover that information because the analysis produces a worse result.

[Mogball]

This makes sense to me and is the right way to handle this.

[Mogball]

The argument states of WarpSpecializeOp should be forwarded from its operands. Unfortunately it doesn't implement RegionBranchOpInterface so you might have to implement this yourself.

[neildhar]

Yeah, I'm planning to look a little more into this, but it seemed orthogonal to this PR. With this PR, we no longer try to populate the state of every operand that hasn't been visited. However, WarpSpecializePartitionsOp is special because we walk the IR and restart the analysis on it. So the framework will call setToEntryState on the block args to it (rather than us calling it ourselves from visitOperation).

When that happens, we want to start with the maximum divisibility, contiguity and constancy, so that we don't start from an artificially restricted state.