Mutually exclusive side effects with a token relationship can cause (avoidable?) scheduling problems

[grebe]

If two side-effecting operations that are ordered by a token relationship have mutually exclusive predicates, the token relationship can be removed.

Removing this token can be useful in scheduling. For example

send0: token = send(tok, state, channel_id=0)
rcv0: (token, bits[32]) = receive(send0, channel_id=1)
rcv0_token: token = tuple_index(rcv0, index=0)
send1: token = send(rcv0_token, send_data, pred=p, channel_id=2)
rcv1: (token, bits[32]) = receive(send1, pred=not_p, channel_id=3)
next_state: bits[32] = tuple_index(rcv1, index=1)

Note that a token relationship exists between send1 and rcv1, but they have mutually exclusive predicates. If rcv0->send1 has a long IO constraint, this IR could be infeasible to schedule because of the long state backedge. Rewriting rcv1: (token, bits[32]) = receive(rcv0_token, pred=not_p, channel_id=3) is legal because send1 and rcv1 have mutually exclusive predicates and makes this IR schedulable.

[ericastor]

This seems more like a case where the token relationship shouldn't be in place? By definition, the optimizer is expected to respect token relationships.

Out of curiosity, can you give a concrete example that produced this IR? This may be more of an issue with the frontend (DSLX or XLS[cc]) failing to let you express the actual ordering relations you want than it is an issue with the IR.

[grebe]

This exact IR was illustrative and didn't come from anywhere. The the situation could come about by a common DSLX style to overspecify token relationships let tok = send(...); let (tok, ...) = receive(...); let tok = send(...) (similar can come from XLScc frontend). It would be better if our token graphs weren't overspecified, but it can be difficult w/ parameterized code.

I think it would be surprising to have a token relationship between side-effects and then end up having the later op scheduled in an earlier stage than the earlier op. Perhaps if two ops are mutex, you can remove anySendThenRecvConstraints? Maybe there could be an optionally-enabled unsafe token graph pruner to try to help the scheduler out? Maybe if scheduling fails, an error message saying "remove this unnecessary edge in the token graph to schedule"?

I've updated the title to not propose a solution and instead point out the issue here.