Both myself and @mattarm have been looking into a more robust way we can partition any directed acyclic graph for different compilers/backends. Below is what we came up with:
What’s wrong with the current partitioning algorithm?

Lack of support for subgraphs with multiple outputs.

If multiple outputs were supported, the current algorithm wouldn’t be able to resolve data dependency issues, meaning that a graph can be partitioned in such a way that forms 2 (or more) uncomputable subgraphs. For example the following graph (with each color representing a different compiler) will be partitioned incorrectly.
The result will be two subgraphs: S1[A  B  D)] and S2[C]. However, this cannot be computed since S1 requires the output from S2 and S2 requires an output from S1. The correct way to partition the graph would be S1[A  B], S2[C], S3[D].
Our proposal
An algorithm which resolves these issues.
The principle behind resolving data dependency issues is the property that a data path which leaves a subgraph cannot reenter that subgraph.
 We start by recursively visiting each node in the input graph  starting at the graphs output nodes and only adding a node to a subgraph once all its parents have subgraphs assigned.
 Once a node is encountered that can be added we create a set of all the subgraphs the node is able to join:
 Check whether any of the parent nodes subgraphs are supported. The issupported nature of each node will be taken care of by the annotation pass makes (this algorithm also opens up the possibility to annotate inline without a separate annotation pass).
 We maintain a ‘restriction map’ which ensures the principle stated above isn’t violated. Once we have a set of supported subgraphs we need to take away the subgraphs that are restricted from one another.
 Following the previous step, we can now add the current node to a subgraph using one of 2 strategies:
 Join an existing subgraph (and then check if a merge is possible)
 Create a new subgraph
 Due to the natural topological ordering of the recursive algorithm, the output is a series of subgraphs that are correctly partitioned.
The restriction map
The restriction map is an unordered map that keeps track of subgraphs as they are added and ensures that once a subgraph has left it (and its parent subgraphs) are no longer rejoined.
 key  the subgraph label that uniquely identifies a subgraph in the graph.
 value  a set containing other subgraph labels that are restricted from joining the subgraph represented by the key.
After a new subgraph is joined or created for a particular node, the restriction map is updated so that the subgraph just assigned is restricted from rejoining any other parent subgraphs plus the inherited restrictions from those subgraphs. Therefore, whenever a path leaves a subgraph, it remembers it’s not allowed to reenter that subgraph by adding the subgraphs it has left as restrictions.
Example
The above much better explained with an example. This particular example shows how the graph can be partitioned correctly involving cases where we: create a new subgraph, join an existing subgraph, merge two subgraphs after joining one and and deal with data dependencies using the restriction map.
Any thoughts are appreciated