Wrong handling of workspace in multiple numpy operators

[ptrendx]

Description

The backward pass of tensordot (both regular and the integer version) operator uses workspace to store intermediate value, but then uses ReduceAxesComputeImpl, which also uses workspace). See e.g. here:
https://github.com/apache/incubator-mxnet/blob/bd55002/src/operator/numpy/np_tensordot_op-inl.h#L420-L428

Since there is only a single workspace storage in MXNet, this means that it is possible for the snippet linked above to

• deallocate dtypespace (if the workspace needs to be reallocated to fit the workspace needed for reduce) leading to a crash
• use the same pointer as input and workspace, leading to silent wrong results

@szha @leezu

[szha]

good catch. it should be refactored to use stateful op instead.

[sxjscience]

@hzfan @reminisce @haojin2 I think we have the discussion about multiple workspaces before. Would you help provide some contexts on this?

[ptrendx]

@szha I don't think one would need stateful op there, just allocate the larger workspace that can hold both the intermediate result and the workspace for reduction.

[ptrendx]

Same problem in MomentsForwardImpl: https://github.com/apache/incubator-mxnet/blob/33d94f1d59335f504ed5b9a7b32f0e81a5d5da56/src/operator/nn/moments-inl.h#L136-L141

[ptrendx]

There is a lot of copy-paste in the implementation of various numpy ops - they should be reviewed. Maybe we could make a test pipeline with a warning if someone tries to access ctx.requested[0].get_space_typed more than once in a given operator and review those usages?

[sxjscience]

I raised the concern before to @reminisce and @haojin2 but there were some discussions that mentioned that we can have multiple workspaces.

[haojin2]

@ptrendx Please refer to the new impl for numpy which uses a pre-allocated workspace to avoid this problem at: https://github.com/apache/incubator-mxnet/blob/master/src/operator/numpy/np_broadcast_reduce_op.h#L666-L694. the MomentsForwardImpl you posted is not serving any numpy operator

[ptrendx]

@haojin2 Sure (I'm actually folding that implementation back to the original for code reuse but anyway), but this function is not used in the examples I showed.

I did see multiple examples of this, including some that I am fixing in my #19426 PR (like this in numpy kron: https://github.com/apache/incubator-mxnet/blob/master/src/operator/numpy/np_kron-inl.h#L229-L234). It just feels that whatever I fix in that PR is just patching things that I noticed instead of making sure the issue is fully fixed. That is the purpose of that issue.