Eliminate common expressions

[ptrendx]

Description

This PR introduces a graph pass that eliminates redundant common expressions in the graph.

For example, let's look at the graph created from a following Python code:

a = mx.sym.Variable()
out = (a + 5) * (a + 5)

This graph will compute a+5 twice, which is wasteful. After the pass introduced in this PR, the actual executed graph will be equivalent to

a = mx.sym.Variable()
b = a + 5
out = b * b

which computes the a+5 only once.

@eric-haibin-lin FYI

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Code is well-documented:
• For new C++ functions in header files, their functionalities and arguments are documented.
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Feature1, tests, (and when applicable, API doc)
• Feature2, tests, (and when applicable, API doc)

Comments

[abhinavs95]

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

[DickJC123]

This will be an awesome addition. Some things to consider in polishing it:

1. The crux is in the definition of functionally equal nodes. Do you think your method bool NodeEqual(const Node * n, const Node * m) belongs in the CSE code, or in the Node class?
2. It's probably best to be conservative with the NodeEqual, but by comparing the attrs.dict, you will miss some equivalent nodes. To illustrate my point, I instrumented the reduce operator to print out the dict:

>>> x = mx.nd.array([1,2,3,4])
>>> mx.nd.sum(x)
[21:32:38] src/operator/tensor/./broadcast_reduce_op.h:621: node attr dict = {}
[10.] <NDArray 1 @cpu(0)>
>>> mx.nd.sum(x, axis=())
[21:32:53] src/operator/tensor/./broadcast_reduce_op.h:621: node attr dict = {{axis,()}}
[10.] <NDArray 1 @cpu(0)>
>>> mx.nd.sum(x, axis=0)
[21:33:07] src/operator/tensor/./broadcast_reduce_op.h:621: node attr dict = {{axis,0}}
[10.] <NDArray 1 @cpu(0)>
>>> mx.nd.sum(x, axis=(0,))
[21:33:39] src/operator/tensor/./broadcast_reduce_op.h:621: node attr dict = {{axis,(0,)}}
[10.] <NDArray 1 @cpu(0)>

All sum operators are functionally equivalent, but none would compare as equal with the approach that includes comparing the attrs.dict (the map<string,string> from the python operator).
Any chance you could move the equality to comparing the parameter struct?
3. I see you consider operators that have resources as never equal. Might you maintain a 'white list' of resources that are OK for equal nodes to have, e.g. the commonly used tempspace resource?
4. I'd prefer to see a test for each of the reasons you deny node equality (e.g. having mutable inputs). You could either count the nodes, to prove the CSE did not happen, or test against a golden copy since performing the CSE would break functionality. I see you have a test for a case where CSE should happen, but it only looks for the reduced node count without testing functionality.
5. Since CSE combines output nodes, those output nodes should not be mutable inputs of downstream nodes, e.g.

x = Variable('x')
xcopy1 = x.copy()
xcopy2 = x.copy()
y = SomeOp(...,some_mutable_input=xcopy1,...)
z = SomeOp(...,some_mutable_input=xcopy2,...)

If you combine equal nodes xcopy1 and xcopy2, then the y and z nodes will start seeing the effects of each other, which they shouldn't as originally coded.

[ptrendx]

Thanks @DickJC123! Your thoroughness is amazing as always :-). I will think about the other points you made, but point 5 is especially interesting. I did not think about it before and you are right that this is a potential problem. I will introduce an additional check for that.

[DickJC123]

Another area to spend some extra thought on involves the primary outputs and the gradients of primary inputs. If you combined two output NDArrays via CSE, would there be any problem with someone wanting to index through the outputs? Ditto for input grads. If a CSE-optimized model generated NDArrays for another model (and that model had mutable inputs), would functionality be preserved? I think a downstream model gets a copy of the NDArray as the input, so it should work, but do we ever want to optimize away that copy?

CSE may eliminate downstream in-place options for node outputs that have been combined. Not sure if there's a case where using CSE is a net perf loss, but I doubt it.

And finally, here's a fairly obscure case where the operator state may not be contained entirely in the node attributes (we should probably outlaw this usage, since it assumes no caching of the env var lookup):

sym = ...
os.environ['MXNET_MY_OP_BEHAVIOR']=X
sym2 = MyOp(data=sym,...)                        # op behavior affected by captured env var setting
os.environ['MNXET_MY_OP_BEHAVIOR']=Y
sym3 = MyOp(data=sym,...)                        # op behavior affected by captured env var setting

Perhaps shifting from the node attr dict to the param struct would catch this case.

[ptrendx]

@DickJC123 Ok, answering your questions:

1. This definition of Node equality is quite tied to what I want to do with them, so I believe, at least right now, it belongs to CSE code.
2. While I agree it would be beneficial, it would also be pretty hard to do - the parameter structure is kept without the type information as the shared_ptr<dmlc::any>. Comparing memory (which could be dine if any gave size of the type held) is also not really feasible since dict is actually part of the Parameter structure, so that would be different.
3. This is the purpose of this THasDeterministicOutput property.
4. I will think about that.

[DickJC123]

I believe the case of having one output driving multiple symbol outputs is handled strangely, not that this is an important use case. The issue occurs with an operator driving 3 or more symbol outputs. Here's the progression of model compute graphs:

foo ----+----> out0
        |
        +----> out1
        |
        +----> out2

copy operators inserted:

foo ----+----> out0
        |
        +-- copy --> out1
        |
        +-- copy --> out2

More CSE applied:

foo ----+----> out0
        |
        +-- copy --+----> out1
                   |
                   +----> out2

copy operators inserted:

foo ----+----> out0
        |
        +-- copy --+----> out1
                   |
                   +-- copy --> out2

So N outputs driven by the same node end up being driven by N - 1 serialized copies, not N-1 parallel copies.

What is the use-case that requires these copies to be inserted?

[DickJC123]

The EliminateCommonExpr pass is performed first-thing, before the backward nodes have been inserted and so before ctx info (i.e. dev_mask) have been added to the graph. I think this is the correct place- doing CSE with backward nodes present would be a nightmare.

I'm not fond of the HasDeterministicOutput flag- seems like it's just one more thing the user has to understand and set correctly when making a new operator. I was hoping that perhaps this property could be inferred by the op's registered resources. Life would be simple if the resources were determined only by the FResourceRequest(attrs) callback. The problem is there's also the FResourceRequestEx(attrs, dev_mask, dispatch_mode) callback, which takes a dev_mask and dispatch_mode. As I point out above, CSE is performed before dev_mask is available from the graph. The only way I can see doing what I'm suggesting would be to infer the possible settings for dev_mask and dispatch_mode from the registration of other callbacks. Some possible pseudo-code for this:

set<resources> all_possible_resources = {}
if (op_has_registered("FResourceRequestEx") {
  auto f_resource_req_ex = < look up op's FResourceRequestEx >
  if (op_has_registered("FCompute<cpu>")
    all_possible_resources += f_resource_req_ex(attrs, kCpu, DispatchMode::kFCompute)
  if (op_has_registered("FComputeEx<cpu>")
    all_possible_resources += f_resource_req_ex(attrs, kCpu, DispatchMode::kFComputeEx)
  if (op_has_registered("FCompute<gpu>")
    all_possible_resources += f_resource_req_ex(attrs, kGpu, DispatchMode::kFCompute)
  if (op_has_registered("FComputeEx<gpu>")
    all_possible_resources += f_resource_req_ex(attrs, kGpu, DispatchMode::kFComputeEx)
} else if op_has_registered("FResourceRequest") {
  auto f_resource_req = < look up op's FResourceRequest >
  all_possible_resources += f_resource_req(attrs)
}
bool inferred_cse_not_possible = LookForBadResources(all_possible_resources)

Regarding the actual name 'THasDeterministicOutput', might this be confused with the cuda notion of deterministic algos? The Convolution op would be tagged with THasDeterministicOutput=1, but would not truly be deterministic if a convolution algo with atomics is used.

[DickJC123]

I've supplied commits to @ptrendx that have this PR to my liking. LGTM.