Ignore non-XLA nodes and their direct dependents.

[ysiraichi]

Fix: #5966

This PR generalizes FallBackNodeCollector into UnsupportedNodesCollector, improving and solving a few issues the old implementation had:

• nodes which resulted in a non-XLA tensor weren't flagged as fallback
• nodes whose arguments were a container with non-XLA tensors weren't flagged as fallback (e.g. stack)
• tracking these "unsupported" nodes weren't really only fallback nodes
  • but, nodes that can't exist in partition boundaries (e.g. arguments and return values)

[ysiraichi]

cc @JackCaoG @miladm

[ysiraichi]

Note: this PR can, possibly, create more partitions than before, affecting performance. Maybe we should make sure there are no regressions before actually landing it.

[JackCaoG]

I think the definition of the fallback is that we have to execute this op on fallback devices(uusually) cpu. for the cases you mentioned

1. nodes which resulted in a non-XLA tensor weren't flagged as fallback
   Does the operation being executed on cpu or gpu/tpu?

2. nodes whose arguments were a container with non-XLA tensors weren't flagged as fallback (e.g. stack)
   I guess this one is technically a fallback, since op will be executed on cpu. It will show up in dynamo fallback messages, but shouldn't be in the XLA's metrics aten counter I think?

I don't fully understand what 3 means, can you give ma an example?

[ysiraichi]

The problem was that FallBackNodeCollector was doing more than just collecting operations that were executed on CPU. It was also flagging nodes, that had non-XLA tensors arguments, as fallback (even though they were executed on XLA).

Reason: my guess is that it did so, in order to guarantee that the input/output of the generated partitions were all XLA tensors (which I believe extract_internal assumes).

Assuming that's the case, it still missed the 2 cases I mentioned:

• nodes which result in non-XLA tensor: might end up as output of some partition
• nodes whose arguments were a container with non-XLA tensors: arguments might end up as the input of some partition

Solution:

• disambiguate fallback from unsupported nodes (from the perspective of the partitioner)
• add the missing 2 cases

@JackCaoG let me know what you think.

[JackCaoG]

yea that make sense, is this pr ready for review?

[ysiraichi]

Yes, it is.

[JackCaoG]

mostly lgtm, minor questions on test cases.

[ysiraichi]

@JackCaoG could you approve this PR?