Memory planner doesn't respect 'output independence'. More optimizations possible. #16685

Description

This issue has two parts. Part 1 involves a discussion of what the copy operator mx.sym.identity() is good for, and how the memory planner currently violates what a user might expect of its behavior. Part 2 is a discussion around more advanced optimizations that the memory planner could be doing.

Part 1. mx.sym.identity() behavior

MXNet users should expect that the output of mx.sym.identity is a copy of its input. Counter to this, the memory planner aggressively combines the storage for the input and output of identity. Basically, the memory planner should be free to make these storage optimizations, as long as they are not visible to the user. See the To Reproduce section below for an example of where the memory planner has gone too far, creating a situation where 2 independent model outputs end up having combined storage (so writes to one output become visible to the other). This issue has relevance to the Common Subexpression Elimination PR #15657, which is being crafted to preserve output independence even as it combines identical functions. The issue of output independence also has relevance to issue #16108 and PR #16131, since the memory planner also runs on subgraphs.

Now the practice of writing-over model NDArrays is not a huge use-case, so it's a shame if MXNet has a memory-footprint or performance penalty due to the desire to preserve output independence. Some possible remedies:

• Have model outputs be read-only (pass-through signals become a problem?)
• Have independent model outputs that have been combined be 'copy on write'
• Have the memory planner be aware of graph i/o 'user visibility', so that it could respect output independence on the main graph, but be more aggressive with internal subgraphs
• Have an argument to bind() specify the policy regarding optimization around I/O storage

Final notes: First, a model can present an input variable directly as an output. In this case, a user can expect that writing to the model output will effect the input. Second, models can present the same operator output multiple times with e.g. mx.sym.group([x, x]). In this case, the framework should preserve model output dependence.

Part 2. More aggressive memory planner optimizations

Currently, the memory planner avoids sharing the storage of primary inputs with any internal node-entries. I believe this policy also effects memory planning of subgraphs. The following optimizations as a result are not performed:

• a primary input passed to an identity op does not share the input's storage with the identity op output
• a primary input passed to an op that has specified in-place but not in-place identity does not share storage with the op's output. This optimization should not be allowed on the top-level graph, as it would introduce the ability of a graph execution to modify its input visibly to the user. However, the optimization might be possible for subgraphs if appropriate graph input node-entry user-visibility and refcount information is provided to the memory planner.

To Reproduce

Run:

def test_output_independence():
    ctx = default_context()
    x = mx.sym.Variable('x')
    x_copy = mx.sym.identity(x)
    y0 = mx.sym.identity(x_copy)
    y1 = mx.sym.identity(x_copy)
    sym = mx.sym.Group([y0, y1])

    x = mx.nd.random.uniform(-1, 1, (4,), ctx=ctx)
    exe = sym.bind(ctx=ctx, args={'x': x}, grad_req='null')
    outs = exe.forward(is_train=True)
    print(' out0 = {}\n out1 = {}'.format(outs[0].asnumpy(), outs[1].asnumpy()))
    outs[0][:] = 0.0
    outs[1][:] = 1.0
    print(' out0 = {}\n out1 = {}'.format(outs[0].asnumpy(), outs[1].asnumpy()))

Which outputs:

 out0 = [ 0.30988264 -0.7911282  -0.92523086  0.85585463]
 out1 = [ 0.30988264 -0.7911282  -0.92523086  0.85585463]
 out0 = [1. 1. 1. 1.]      # <- wrong!!
 out1 = [1. 1. 1. 1.]

@ptrendx @szha @eric-haibin-lin @samskalicky