Fix for duplicate subgraph inputs/outputs

[samskalicky]

Description

Fixes #16108

Subgraph API can create duplicated tensors for both inputs to a subgraph and outputs from a subgraph. This happens when there are multiple nodes in a subgraph that consume the same input, or when there are multiple nodes that consume a single subgraph output. This results in tensor duplication, and causes OOM due to excessive memory usage.

Fix for duplicate inputs

Currently, inputs to a subgraph are duplicated when there are multiple nodes in a subgraph consuming the same input (see diagram below). In this PR we change the behavior of CutGraphInputs function to store a NodeEntry for each input to the subgraph in a map, and re-use that node for each node in the subgraph that consumes that input (see diagram below). This prevents input nodes from being duplicated, and results in less copies of tensors at runtime and a much smaller memory footprint.

Fix for duplicate outputs

Similar to the inputs, when an output from a subgraph is consumed by multiple nodes outside the subgraph, the output is duplicated for each consumer (see diagram below). In this PR we change the behavior of CreateSubgraphNode to collapse duplicate output dependencies to a single output by checking if neighboring output entries are the same. If so we do not add another output, preventing duplicates.

Then in the ConnectSubgraphOutputs function we create dependencies from each node outside the subgraph that consumes a subgraph output. If neighboring nodes share the same subgraph output, they point to the same subgraph output.

Example

Added an example graph with duplicate inputs/outputs to the test_subgraph_op.py. Heres the original graph partitioned before this PR. Notice that it has two input nodes: data0 and data1 that both come from the same top level data. Similarly, notice that it has two outputs coming from the same node _plus0.

{
  "nodes": [
    {
      "op": "null", 
      "name": "data", 
      "attrs": {"__shape__": "(1,)"}, 
      "inputs": []
    }, 
    {
      "op": "_CachedOp", 
      "name": "_CachedOp0", 
      "inputs": [[0, 0, 0], [0, 0, 0]], 
      "subgraphs": [
        {
          "nodes": [
            {
              "op": "null", 
              "name": "data0", 
              "inputs": []
            }, 
            {
              "op": "sin", 
              "name": "sin0", 
              "inputs": [[0, 0, 0]]
            }, 
            {
              "op": "null", 
              "name": "data1", 
              "inputs": []
            }, 
            {
              "op": "sin", 
              "name": "sin1", 
              "inputs": [[2, 0, 0]]
            }, 
            {
              "op": "elemwise_add", 
              "name": "_plus0", 
              "inputs": [[1, 0, 0], [3, 0, 0]]
            }
          ], 
          "arg_nodes": [0, 2], 
          "node_row_ptr": [0, 1, 2, 3, 4, 5], 
          "heads": [[4, 0, 0], [4, 0, 0]]
        }
      ]
    }, 
    {
      "op": "cos", 
      "name": "cos0", 
      "inputs": [[1, 0, 0]]
    }, 
    {
      "op": "cos", 
      "name": "cos1", 
      "inputs": [[1, 1, 0]]
    }, 
    {
      "op": "elemwise_sub", 
      "name": "_minus0", 
      "inputs": [[2, 0, 0], [3, 0, 0]]
    }
  ], 
  "arg_nodes": [0], 
  "node_row_ptr": [0, 1, 3, 4, 5, 6], 
  "heads": [[4, 0, 0]], 
  "attrs": {"mxnet_version": ["int", 10700]}
}

Heres the same graph partitioned after this PR. Notice that now there is only 1 input node data0 and both sin ops use it. Also notice that now there is only 1 output from the subgraph, and both cos ops use it.

{
  "nodes": [
    {
      "op": "null", 
      "name": "data", 
      "attrs": {
        "__profiler_scope__": "<unk>:", 
        "__shape__": "(1,)"
      }, 
      "inputs": []
    }, 
    {
      "op": "_CachedOp", 
      "name": "_CachedOp0", 
      "inputs": [[0, 0, 0]], 
      "subgraphs": [
        {
          "nodes": [
            {
              "op": "null", 
              "name": "data0", 
              "inputs": []
            }, 
            {
              "op": "sin", 
              "name": "sin0", 
              "attrs": {"__profiler_scope__": "<unk>:"}, 
              "inputs": [[0, 0, 0]]
            }, 
            {
              "op": "sin", 
              "name": "sin1", 
              "attrs": {"__profiler_scope__": "<unk>:"}, 
              "inputs": [[0, 0, 0]]
            }, 
            {
              "op": "elemwise_add", 
              "name": "_plus0", 
              "attrs": {"__profiler_scope__": "<unk>:"}, 
              "inputs": [[1, 0, 0], [2, 0, 0]]
            }
          ], 
          "arg_nodes": [0], 
          "node_row_ptr": [0, 1, 2, 3, 4], 
          "heads": [[3, 0, 0]]
        }
      ]
    }, 
    {
      "op": "cos", 
      "name": "cos0", 
      "attrs": {"__profiler_scope__": "<unk>:"}, 
      "inputs": [[1, 0, 0]]
    }, 
    {
      "op": "cos", 
      "name": "cos1", 
      "attrs": {"__profiler_scope__": "<unk>:"}, 
      "inputs": [[1, 0, 0]]
    }, 
    {
      "op": "elemwise_sub", 
      "name": "_minus0", 
      "attrs": {"__profiler_scope__": "<unk>:"}, 
      "inputs": [[2, 0, 0], [3, 0, 0]]
    }
  ], 
  "arg_nodes": [0], 
  "node_row_ptr": [0, 1, 2, 3, 4, 5], 
  "heads": [[4, 0, 0]], 
  "attrs": {"mxnet_version": ["int", 20000]}
}

[samskalicky]

@zheng-da @reminisce could you please review?

[samskalicky]

@mxnet-label-bot add [pr-awaiting-review]

[samskalicky]

@ZhennanQin for review. Any suggestions on how to better fix the problem?

[samskalicky]

@ZhennanQin ive made the changes, do they still look good? If so we can start going through and updating all other subgraph properties.

[ZhennanQin]

LGTM. Please coordinate all backends with new ConnectSubgraphOutputs().

[DickJC123]

I've noticed that MXNet does not preserve what I'm calling 'output independence' during it's memory planning (issue forthcoming). While I'm talking here about the model as a whole, I've noticed many graph passes are applied to subgraphs as well, so perhaps the issue is relevent to this PR. Any update on status here?

[samskalicky]

Hi @DickJC123 sorry I havent been able to prioritize finishing off this PR. IMO the problem is similar and endemic in mxnet but not related to this PR. Please do file an issue at the very least for posterity.

[samskalicky]

@mseth10 looks like we made an aux an arg somehow, thats why the MKLDNN tests are failing:

in_args_map	arg_names	arg_names
data	data	data
conv1_weight	conv1_weight	conv1_weight
bn1_gamma	bn1_gamma	bn1_gamma
bn1_beta	bn1_beta	bn1_beta
conv2_weight	conv2_weight	conv2_weight
bn2_gamma	bn2_gamma	bn2_gamma
bn2_beta	bn2_beta	bn2_beta
		bn2_moving_mean

So this is why the error says

MXNetError: Check failed: arg_names.size() == in_args_map.size() (8 vs. 7) : 

Theres one extra name in arg_names for bn2_moving_mean

[HahTK]

There is an alternative implementation that does not need t track the actual unique inputs separately.
A counter would allow us to correctly modify orig_entries into unique_entries.
This keeps the function signature unchanged and minimizes changes elsewhere.

It is already working in a private build. I add it here for consideration if desired

[samskalicky]

Thanks for the suggestion @HahTK, but we cant modify orig_entries since we need it unmodified for ReattachGraphInputs:
https://github.com/apache/incubator-mxnet/blob/e249d71ad4621afcba7f0f3af77095a3f9a4bc83/src/operator/subgraph/build_subgraph.cc#L578
this is used to reject a subgraph from the reviewSubgraph API:
https://github.com/apache/incubator-mxnet/blob/e249d71ad4621afcba7f0f3af77095a3f9a4bc83/src/operator/subgraph/build_subgraph.cc#L657
the example you provided doesnt work in the latest version of MXNet where you can reject a subgraph by returning nullptr from CreateSubgraphNode here: https://github.com/apache/incubator-mxnet/blob/e249d71ad4621afcba7f0f3af77095a3f9a4bc83/src/operator/subgraph/build_subgraph.cc#L630-L633

[HahTK]

The proposed change was tested on 1.5.1 and it seems like newer versions of MXNet have new features that needs orig_entries to be preserved. So acknowledged !

[waytrue17]

Looks good to me, thanks!

[mseth10]

LGTM. Thanks for the fix.

[samskalicky]

@mseth10 looks like we made an aux an arg somehow, thats why the MKLDNN tests are failing:

in_args_map arg_names arg_names
data data data
conv1_weight conv1_weight conv1_weight
bn1_gamma bn1_gamma bn1_gamma
bn1_beta bn1_beta bn1_beta
conv2_weight conv2_weight conv2_weight
bn2_gamma bn2_gamma bn2_gamma
bn2_beta bn2_beta bn2_beta
bn2_moving_mean
So this is why the error says

MXNetError: Check failed: arg_names.size() == in_args_map.size() (8 vs. 7) : 

Theres one extra name in arg_names for bn2_moving_mean

For the record @mseth10 i finally found the issue that caused this "making an aux an arg" problem: #19112 (comment)