in-place reshape ops

[szha]

Description

make NDArray fluent methods for expand_dims/flatten/squeeze in-place reshapes.

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
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• Check the API doc at http://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• make NDArray fluent methods for expand_dims/flatten/squeeze in-place reshapes.

Comments

• resolves expand_dims() makes copy instead of simply reshaping #13998

[ZhennanQin]

Have you tried this with mkldnn enabled? E.g reshape the output of mkldnn convolution?

[szha]

@ZhennanQin I didn't test that specifically based on the assumption that the backend implementation should not cause existing APIs such as NDArray.reshape or NDArray.shape to break.

[stephenrawls]

@szha -- Thanks for putting in this patch!

I have a couple questions about the operator expand_dims, which I understand you are actually changing the python code to not use that operator, but I think still has relevance.

(1) I originally discovered this because I was using the C API to call MXImperativeInvoke() using the expand_dims operator, and I noticed it was causing a copy. This fix only effects the Python version when operating on ndarrays, not users of the expand_dims operator.

I see that there is a path in the operator that checks if it is an in-place operation. Presumably it uses that path if I pass an output array that is the same NDArrayHandle as the input array? But what if I still need the original input array handle, and I want to create a new output array handle with the expanded dim but still not make a copy?

(2) In the issue I created you commented that: "For symbol (and thus the hybridized version), since in-place identity is possible it should not matter".

Can you talk a little more about that? I assume you mean that in this case:

x_expanded = x.expand_dims(1)
y = x_expanded + foo

The engine can figure out that x is not needed again, and can thus turn the expand_dims(1) into an in-place operation that doesn't make a copy?

I'm not very familiar with how this part of the code works, so what happens if you had code that looked like this?

x_expanded = x.expand_dims(1)
y = x_expanded + foo
z = 2 * x

i.e. the code still makes a reference to the original x, and thus presumably the engine can't decide to use the in-place version of expand_dims in that case, right? So I guess my question is -- Does the ability for the Syblolic / hybridized engine to elide the copy depend on the code not referencing the un-expanded version of the array after calling expand_dims()? If so, it seems like there will still be some use cases where an unexpected copy is happening.

[szha]

@stephenrawls

This fix only effects the Python version when operating on ndarrays, not users of the expand_dims operator.

Right. The reshape operators always return a copy. This is because in imperative mode there is no complete computation graph to be analyzed and see whether a copy is needed or not.

I see that there is a path in the operator that checks if it is an in-place operation. Presumably it uses that path if I pass an output array that is the same NDArrayHandle as the input array?

Right, it is done through output=original_array argument to the operators. It does not work for autograd (i.e. Gluon's training mode) as autograd doesn't support in-place operation.

But what if I still need the original input array handle, and I want to create a new output array handle with the expanded dim but still not make a copy?

As long as it's not used in the context of autograd, you can save the original handle before in-place operation. Though the output of in-place op shares the space with the input, it's still a separate handle.

For symbol (and thus the hybridized version), since in-place identity is possible it should not matter.

To elaborate, in symbolic mode, the backend can see the complete computation graph, and as a result it can analyze the graph (through node coloring) and figure out whether nodes can share the same space. FInplaceIdentity is the graph node attribute that expresses whether this should be allowed in order for compiler to plan memory. For details see 3rdparty/tvm/nnvm/src/pass/plan_memory.cc.

Note that such memory plan happens before the engine/scheduler work happens, and it happens for both symbolic executor and hybridized Gluon HybridBlock.

[stephenrawls]

@szha -- Thanks for the detailed explanation, that helps a lot.

So just for clarity, I see that in the operator definition for expand_dims, it sets the FInPlaceIdentity attribute to true:
https://github.com/apache/incubator-mxnet/blob/45d1a1e6ff9c58cfc75f72651c1bf671ac7f1885/src/operator/tensor/matrix_op.cc#L400-L419

I didn't see anything in the plan_memory.cc file that depended on autograd / training.

So should I take it that the code is smart enough to not allocate new memory for the expand_dims operator in symbolic mode, i.e. not to make a copy, even during training when using autograd?

If so, then I think I understand it, and I see what you mean about only needing to fix this for the Python ndarray use case.

Thanks again!
Stephen

[vandanavk]

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

[szha]

@stephenrawls actually, autograd comes into play in pure imperative mode. MXNet's autograd is trace based and does not support in-place operation such as input = mx.nd.expand_dims(input, ..., output=input). You're actually right in that it doesn't help the case where you write it in imperative mode in C++. In order to make use of the graph pass for saving memory, you need to use Symbol API to express the model and use symbolic execution (or its counter part for Gluon, which is CachedOp).

[TaoLv]

Does it mean that with these changes python API will have different behavior compared with other frontend languages?

[szha]

@TaoLv Other language bindings don't seem to have the fluent methods for reshaping NDArray.

[szha]

The current implementation would break backward compatibility for in-place update.

[szha]

I plan on adding an inplace argument to these fluent methods to allow users to specify whether in-place operation should be used.

[ankkhedia]

@reminisce Could you please help review the PR?

[szha]

I haven't finished this PR yet as I haven't got around to add the in-place option.

[ankkhedia]

@mxnet-label-bot add [pr-work-in-progress]

[ankkhedia]

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

[reminisce]

Please add unit tests.

[szha]

@reminisce good catch. I was relying on existing tests before adding the inplace option but forgot to add tests for them afterwards.

[szha]

Added tests

[wkcn]

LGTM : )

I have a question.

It seems that in-place flag does not support Symbol.
In gluon, F.expand_dim(a, inplace=True) will raise inplace not found when hybridizing the model.

Could we add warnings in mxnet/symbol.py ?

[szha]

@wkcn good point on symbol and HybridBlock. I'm not sure if a warning should be added as it can get very verbose.

[szha]

Nonetheless it should not throw an error when using inplace with hybridize. Let me see how to best deal with it.

[karan6181]

@szha Thank you for your contributions! Is this PR work in progress?

[abhinavs95]

@szha is this PR ready to merge?

[piyushghai]

@szha Is this good to merge ?

[Roshrini]

LGTM

[szha]

Thanks for all the reviews. I plan to add a dummy option in the symbol version of the reshape ops so that it won't throw error when people use this option in the context of Gluon

[roywei]

@mxnet-label-bot add[pr-work-in-progress]

[roywei]

@mxnet-label-bot remove [pr-awaiting-response]

[karan6181]

@szha Did you get a chance to add a dummy option in the symbol version? Thanks!

[piyushghai]

@szha @wkcn This PR is ready to be merged I guess ?

@mxnet-label-bot Update [NDArray, pr-awaiting-merge]

[szha]

@piyushghai thanks for pinging. I wanted to hold onto this for the 1.5 release code freeze.

[abhinavs95]

@szha is this good to go now?