[Discuss] Annotation Defined Subgraphs

Development pre-RFC
#1

In the new external codegen infrastructure, we make significant use of subgraphs denoted by introducing annotations into the relay expressions. The principle is that you can define a subgraph by introducing an annotation op on all of its incoming and outgoing ‘edges’. This is currently used in the AnnotateTarget pass and the PartitionGraph pass to designate subgraphs as being possible to offload to an external codegen (compiler_begin and compiler_end annotations are used).

Ultimately these annotated subgraphs are transformed into functions marked as external by PartitionGraph. However, currently this only works for subgraphs which have a single output. Going forward, we will need to be able to support subgraphs with potentially many outputs to allow us to partition the graph more extensively.

We intend to implement the algorithm described here to support multiple outputs, but this will require the addition of a few extra passes. Approximately speaking, we have three steps:

  1. Annotate operators (or composite functions) as supported by introducing compiler_begin/compiler_end annotations around them.
  2. Take these supported operators and merge them together into larger subgraphs which can be offloaded (the algorithm is described in the aforementioned RFC).
  3. Take the resulting merged subgraphs and turn them into external functions (performed by the PartitionGraph pass).

In all of these the notion of annotated subgraphs is used, so we propose a new bit of infrastructure to extract and manipulate these subgraphs - SubgraphSet (name subject to change…)

A SubgraphSet is initialised using a Relay expression and the subgraph_begin/subgraph_end annotation types (compiler_begin and compiler_end in our case). It will then populate itself with a number of Subgraph objects with information on the nodes they contain as well as the input and output nodes of the Subgraph. It also provides methods to query/manipulate the created Subgraphs.

Having this as a separate utility is valuable as it can then be reused in both the merging and partitioning pass (as well as any future passes that need to work with Subgraphs).

There is a WIP PR here: https://github.com/apache/incubator-tvm/pull/5030.

This is intended as the first step towards supporting multiple outputs in the partitioner. The next step will be an update to PartitionGraph to implement step 3, then followed by a new pass to implement step 2. Concurrently, we will need to add support for composite functions.

@zhiics @comaniac

#2

The overall idea of maintaining a set of subgraph utilities looks good to me. Some high-level comments:

  • The naming will indeed be a problem (sigh). Perhaps you should open another topic with a poll of propose names.

  • I took a brief look at the WIP PR and summarized the APIs to two categories:

    1. Subgraph maintenance, including create, remove, lookup, etc.
    2. Subgraph manipulation, including add node, merging, etc.

While it’s intuitive to include the first category to SubgraphSet, the second category seems a bit weird, because those functions are static. Maybe we can escalate “Subgraph” to a class and put the second categories in it so that it would be like subgraph->Merge(that) or subgraph->AddNode(expr)? In this way, we may even use the existed Map container to maintain subgraphs.

#3

Some high level thoughts:

Terminogolgy wise, it would be great if we look beyond the notion of graph, given that “graph” is not well defined once we move beyond the dataflow only programs.

Can we instead think of alternative name(should they be functions?, or basic block given the lack of control flow?). It is also useful to think about what containers and features we need.

For example, can we use existing containers like IRModule? What are additional features we need for the container(e.g. de-duplication?)

Finally, if we are exposing the comon util as a part of the python API, it would be useful to think about namspacing. e.g. should we put it under relay.analysis?

#4

@comaniac I think you’re right about the responsibilities of SubgraphSet being a bit unclear. My reasoning was that for things like merge we need to remove one of the subgraph from the set. I can probably move the actual merging logic to a method of Subgraph though and just have SubgraphSet manage the erasing of one of the graphs.

@tqchen I anticipated there being some discussion on terminology :slight_smile: I’m not particularly convinced by Subgraph myself. The problem with representing these as functions is primarily because of the multiple outputs scenario. Additionally, they need to faithfully capture the data flow of the program so that we can determine whether Subgraphs can be validly partitioned. This means we need to put a subgraph_end annotation before every consumer of the Subgraph.

One thing I need to think about is whether it will ever be meaningful to merge subgraphs across scopes. Intuitively I think probably not, in which case the ‘subgraphs’ would have to live inside a basic block.

I don’t have a strong view on the namespacing (I just saw CallGraph recently went under relay.analysis so copied that). If you think there’s a more appropriate place, we can move it.

#5

The multiple output case can be represented hopefully as a tuple, so in that case they can be represented as a Function. relay.analysis sounds like a good place to start, perhaps we should bring analysis as a subfolder as we start to grow more features on that namespace

#6

We can use a tuple, but this requires mutating the expression in a ‘destructive’ way. Ideally, we should be able to annotate many different subgraphs that potentially overlap (corresponding to what various targets support) which can’t be easily handled with a function based approach. Turning them into functions is what is currently done in the PartitionGraph pass, and we do this once we’ve finalised what target each subgraph should go to.

#7

This is a bit similar to the recently proposed ExtractFusedFunction RFC, but of course after different passes. I think it should be part of analysis as well.

I think the problem of multiple output was because we assume the last one is the output. Now since we have ADTObj, these problem should be solved by using tuple and packing the runtime object into ADTObj.

@tqchen Yes, I also wanted to separate the passes into transform and analysis subfolders. There are already enough analysis passes now. I believe we will have more once the pattern infra and basic block normal form are introduced. LLVM also has analysis and transforms folders.

#8

One thing I should make clearer is this is a pure annotation mechanism, and the annotations are not necessarily going to be disjoint (eg. in our case two targets could both annotate the same nodes as supported). This is a utility to be able to easily query how an expression has been annotated and we refer to those annotated regions as ‘subgraphs’. We’re constructing these ‘subgraphs’ somewhat speculatively, there will be specific passes whose job it is to deconflict annotations, merge them and then replace them with (global) functions.

I think because of this it may not be appropriate to use a container like IRModule.

Some potential names under consideration include:

AnnotatedNodes
AnnotatedBlock
(Annotated)NodeGroup
AnnotatedRegion
AnnotatedFlowRegion

To further clarify, it’s definitely our intent that once we have confirmed which parts of the module are being offloaded to which codegen the annotations will all be removed and the ‘subgraphs’ will be converted into functions returning a tuple. That just happens to be a later step in the process.

[RFC][BYOC] An extended graph partitioning flow
#9

Annotated sgtm, it seems to be more like an internal analysis tool rather than user facing API. Which might makes it easier to push for now.

The main reason I bought up the naming issue is that we want to minimize mental burden for our developers when dealing with major concepts. so far they are

  • Pass(transformation) and IRModule(the value being transformed)