Fix and optimize handling of vectorized memory accesses

[ptrendx]

Description

For operators, which performance is limited by global memory bandwidth, it is important to issue the widest possible loads, as it ensures that the bandwidth is fully utilized.

Currently in MXNet we use vectorized loads and stores only for half_t type for a few operators (some elementwise binary operators and elementwise_sum). Unfortunately, the way it was done makes some assumptions about MXNet's NDArrays which do not hold true for all cases.

• Failure 1:

import mxnet as mx
ctx=mx.gpu()
a = mx.nd.array([1,2,3,4], dtype='float16', ctx=ctx)
b = mx.nd.array([1,2,3,4], dtype='float16', ctx=ctx)

c = a[1:3]
d = b[1:3]
mx.nd.elemwise_add(c, d, out=c)

Results in error:

Check failed: e == cudaSuccess: CUDA: misaligned address

• Failure 2:

import mxnet as mx
ctx=mx.gpu()
a = mx.nd.array([1,2,3,4], dtype='float16', ctx=ctx)
b = mx.nd.array([1,2,3,4], dtype='float16', ctx=ctx)

print(a)
c = a[0:3]
d = b[0:3]
mx.nd.elemwise_add(c, d, out=c)
mx.nd.waitall()
print(c)
print(a)

gives:

[1. 2. 3. 4.]
<NDArray 4 @gpu(0)>

[2. 4. 6.]
<NDArray 3 @gpu(0)>

[2. 4. 6. 8.]
<NDArray 4 @gpu(0)>

which is a silent data corruption (the last element of a should not have been changed).

It was not noticed before because a + b for NDArrays launches the broadcast_add instead of elemwise_add (and is not vectorized), whereas in the symbolic execution slices give new allocations, which do not exhibit those issues.

This PR:

• fixes those issues
• introduces helpers for handling vectorization (for all types, not only half_t)
• increases performance of vectorized kernels
• introduces vectorization for all binary/unary/binary with scalar ops
• (WIP) introduces vectorization for broadcast ops

@eric-haibin-lin @sxjscience @haojin2

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 best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• Properly handle vectorized loads and stores in elementwise kernels
• Handle vectorized loads and stores in elementwise broadcast kernels

[leezu]

Tangential question related to MSHADOW_XINLINE use in this file: Do we still need inline __attribute__((always_inline)) or is nvcc smart enough to inline? If so, can we stop using the mshadow macros and just use __device__ __host__?

[ptrendx]

MSHADOW_XINLINE is useful because it is empty if you are not running with NVCC (and GCC does not understand __host__ __device__).

[leezu]

True. MSHADOW_XINLINE further sets __attribute__((always_inline)). Do we need that for nvcc?

With respect to __host__ __device__, this file is wrapped into MXNET_USE_CUDA && __CUDACC__, so gcc won't see it.

[ptrendx]

Frankly, not sure - generally speaking inline should be enough, but it is only an advice and compiler is allowed to not inline for whatever reason.

[ptrendx]

@haojin2 I'm hitting again the OoM on Windows build, like you did before. I started looking at the numpy versions of those functions and I see that you have way more templates there (some of which are actually not compiled on Windows) and I started thinking that we should probably just switch to runtime compilation of those kernels - there is just too many variants here. What do you think about this (also @eric-haibin-lin @szha @leezu for comments)?

Also - I don't see elementwise ops in the numpy python package, just broadcast ops - this is pretty bad because knowledge that the shapes are the same is pretty important in optimizations - for example the pointwise fusion would not really work for such operators.

[leezu]

For OoM, could updating the windows toolchain help #17808 ? cc @vexilligera @josephevans

[leezu]

@mxnet-bot run ci [windows-cpu]

[mxnet-bot]

Jenkins CI successfully triggered : [windows-cpu]

[DickJC123]

One thing I like to see, in the spirit of test-driven development, is a first commit that includes the test showing the problem fixed, if possible. Since this was not done, I verified that both sub-tests of test_elementwise_ops_on_misaligned_input cause that test to fail on a gpu (i.e. as invoked via test_operator_gpu.py) on a checkout that doesn't include this PR. The other 2 supplied tests do not fail on this prior checkout, as they were designed to validate the broadcast vectorization that will be part of the follow-on PR. All 3 tests pass when run on the CPU, since this is a GPU issue being addressed here, but good to keep the test coverage consistent.

LGTM.