[Relay] Generate sub graphs from existing graph

popojames · March 1, 2022, 11:37pm

Thanks for raising this post, I am also interested in generating some subgraphs from an existing graph to run on different CPU/accelerators.

In my previous work, I have followed @hjiang’s old post to split the existing graph into N different subgraphs.

github.com

huajsj/tvm/blob/47c5cc24dc01248b0c1b7ea76cb3ff2806445888/tests/python/relay/test_pipeline_executor.py#L25


# specific language governing permissions and limitations
# under the License.


import numpy as np
import tvm
import tvm.testing
from tvm import relay
from tvm.relay import transform
from tvm.contrib import graph_executor, pipeline_executor


"""
Split graph into a serial of sbgraph.
"""
def pipeline_graph(expr, indices):
    """Split Graph Into A Group Of Subgraph
    Parameters
    ----------
    expr : tvm.relay.Expr
    indices : Array[int]
    Returns
    -------

However, as my previous post mentioned, I found out each subgraph can only have one global output, wherein is the last operation.

When I check the data dependency, I notice there are another dependencies other than the last operation:

For example, in my post, %42 of the first subgraph → %x1: Tensor[(1, 1, 1, 128), float32] of the second subgraph. This operation is constant that goes to every layer. (e.g, %19 in second subgraph) However, for this operation, I cannot send the data dependency to the next subgraph since it is not registered as global output in the first subgraph.

Thus, I am wondering is it possible for that user can we register operations in Relay IR as new outputs to read them out (or send them to another subgraph, in my case).

Moreover, can @abhikran-quic share more information regarding which documents do you follow or how do you use the PartitionGraph function?

Thanks for your help

cc @comaniac