[MXNET-1185] Support large array in several operators (part 1)

[apeforest]

Description

This PR fixed the large array issue (#13036, #13070) in the following operators:
ndarray.ones
ndarray.zeros
ndarray.sum
ndarray.slice
ndarray.random.uniform
ndarray.empty

This PR is only the first effort that addresses some limitations in some basic operators. More tests and fix are coming as we identify more operators. This is part of a large project tracked in JIRA by https://issues.apache.org/jira/browse/MXNET-1184

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Nightly tests are added for complicated/long-running ones
• Code is well-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

• Support large array (> 3billion elements) in certain operators
• Added nightly tests on arrays of 5 billion elements

Comments

This PR only supports the total number of elements in the array to go beyond 2 billion. The max shape size in any of the dimension is still limited to 2^31 because this PR does not update the C API. Due to the large change required in all language bindings, we decided to break the original PR (#13191) into several smaller parts. And this PR is part 1.

This PR only changes the indexing variable type from int32_t to int64_t, not the data type of each element. No performance impact observed. For future guidance of data types to choose, please refer to https://cwiki.apache.org/confluence/display/MXNET/Large+Tensor+Support

[zheng-da]

The PR looks good to me.

[zheng-da]

thanks for the great work. your code modifies a lot of code and the test doesn't cover all of them.
could you please add more unit tests to make sure the modified code is covered?
Here I tried a bunch of tests. Most of them work. Unfortunately, broadcast div fails. Could you please double check the problem in broadcast div?

import mxnet as mx
import numpy as np
long_dim = 50000000
arr = mx.nd.ones(shape=(long_dim, 200))

deg = mx.nd.ones(shape=(long_dim,1)) * 2
# Test broadcast div
res = arr/deg
assert np.sum(arr[-1].asnumpy() == 0.5) == arr.shape[1]

# Test element-wise
arr2 = mx.nd.ones(shape=(long_dim, 200))
res = arr + arr2
assert np.sum(res[-1].asnumpy() == 2) == arr.shape[1]
res = arr + 1
assert np.sum(res[-1].asnumpy() == 2) == arr.shape[1]
res = mx.nd.sqrt(arr + 3)
assert np.sum(res[-1].asnumpy() == 2) == arr.shape[1]

# Test reduce
assert mx.nd.sum(arr).asnumpy() == arr.shape[0] * arr.shape[1]

# Test dot
weight = mx.nd.ones(shape=(200, 100))
res = mx.nd.dot(arr, weight)
assert np.sum(res[-1].asnumpy() == 200) == weight.shape[1]

# Test FullyConnected
res = mx.nd.FullyConnected(arr, weight, num_hidden=weight.shape[1], no_bias=True)
assert np.sum(res[-1].asnumpy() == 200) == weight.shape[1]

# Test broadcast
range = mx.nd.arange(0, long_dim).reshape(long_dim, 1)
res = mx.nd.broadcast_to(range, shape=(range.shape[0], 200))
assert np.sum(res[-1].asnumpy() == long_dim) == res.shape[1]
res = mx.nd.broadcast_like(range, arr)
assert np.sum(res[-1].asnumpy() == long_dim) == arr.shape[1]

# Test clip
data = res
res = mx.nd.clip(data, a_min=100, a_max=1000)
assert np.sum(res[-1].asnumpy() == 1000) == arr.shape[1]


# Test take
idx = mx.nd.arange(long_dim-1000, long_dim)
res = mx.nd.take(arr, idx)
assert np.sum(res[-1].asnumpy() == 1) == res.shape[1]

# Test slice
res = mx.nd.slice(arr, begin=(long_dim-1000, 1), end=(long_dim, 100))
assert np.sum(res[-1].asnumpy() == 1) == res.shape[1]

# Test slice assign
res = arr.copy()
res[long_dim-1:long_dim] = 1000
assert np.sum(res[-1].asnumpy() == 1000) == arr.shape[1]

# Test expand_dims
res = mx.nd.expand_dims(arr, axis=1)
assert res.shape == (arr.shape[0], 1, arr.shape[1])

# Test squeeze
data = res
res = mx.nd.squeeze(data)
assert sum(res.shape == arr.shape) == 2

[vandanavk]

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

[ChaiBapchya]

This is pretty useful. However, going forward to ensure consistency, could you guide us as to
when to use
size_t, index_t and auto (instead of the usual int)
Thanks!

[samskalicky]

LGTM. How are the test coming? If you need help with the test creation, can you point out a group of things needing testing and i'll help out.

[apeforest]

@zheng-da I have fixed the bug in broadcast_div and added the tests you provided. Thanks for your review and testing!

[zheng-da]

Could you fix this bug as well? thanks.

def run_linear():
    data = mx.nd.ones(shape=(50*1000*1000, 100))
    linear = gluon.nn.Dense(100)
    linear.initialize(ctx=mx.cpu(0))
    res = linear(data)
    res.wait_to_read()

[apeforest]

@zheng-da I don't see problem with this test. I have added it to the nightly test script. Please review.

[apeforest]

@ChaiBapchya Yes, I will document this in cwiki after this change is merged. In general, index_t should be used for indexing elements in the tensor. size_t is used for returning the size of object or total number of elements. auto is a keyword in C++11 which derives the data type based on the rhs.

I found the Google C++ style guide to be a very good one to follow in programming: https://google.github.io/styleguide/cppguide.html#Integer_Types

[zheng-da]

@apeforest i can't run the test for the dense layer in my C5.18x.

[zheng-da]

do you build your code with mkldnn or without mkldnn?

[zheng-da]

i can verify that the test code for dense works when mxnet isn't built with mkldnn.

[anirudh2290]

@pengzhao-intel raised some good concerns on #13036 about performance. Also from @wkcn tests I see div op performance affected between int32 and int64. Have we done any micro-benchmarks for the operators affected ?

[lebeg]

👍

[apeforest]

@anirudh2290 This PR only changes the indexing variable type from int32_t to int64_t, not the data type of each element. So the performance impact is different from what @wkcn test.

I also did a test in broadcast_div operator in MXNet and found no obvious runtime change (int64_t version actually is slightly faster)

import mxnet as mx
import numpy as np
from mxnet.test_utils import check_speed

a = mx.sym.ones(shape=(100000, 10000))
b = mx.sym.ones(shape=(100000, 1)) * 2
c = mx.sym.broadcast_div(a, b)

print(check_speed(c))

Master branch: 15.357867193222045 seconds
PR branch: 12.930877685546875 seconds

[yuxihu]

LGTM. You have some good information in the conversation which I suggest to add to the comments section of the PR. It will make things clearer and easier for reviewers to catch up in the first place. I think you should mention: (1) This PR only changes the indexing variable type from int32_t to int64_t, not the data type of each element. No performance impact observed. (2) Provide a cwiki link where you will add guidance like: "In general, index_t should be used for indexing elements in the tensor. size_t is used for returning the size of object or total number of elements. auto is a keyword in C++11 which derives the data type based on the rhs."

[apeforest]

@yuxihu Thanks for your nice suggestion. I have updated the PR comment section.

[anirudh2290]

Thanks for the clarification. Looks good overall. Can we also add a unit test for scatter_nd.

[apeforest]

@anirudh2290 Given the overhead of repeating CI, can we get this merged and add unit test for scatter_nd in an second PR? This PR is only part-1 of a larger PR anyways. Thanks.

[pengzhao-intel]

@pengzhao-intel raised some good concerns on #13036 about performance. Also from @wkcn tests I see div op performance affected between int32 and int64. Have we done any micro-benchmarks for the operators affected ?

@anirudh2290 Agree with you.
I see the changes make the INT64 by default in several OPs.
Is it a little aggresive? How about use template to switch the data type?

[zheng-da]

@pengzhao-intel I think the performance impact is minimal. Only the index type is changed. The operations on the index is basically add and comparison. All other computation is much more expensive. I double there is measurable performance difference.

[pengzhao-intel]

@zheng-da @apeforest I am fine with the current PR since it's only the first part of the changes.

I have re-visited the proposal in here and I suggest to add more details and data, such as how CUDA and MKL-DNN supports INT64 type and any potential impacts.

[anirudh2290]

Please add a TODO (JIRA) for the remaining test. Also, makes sense to bring this proposal up for discussion on the dev list.

[apeforest]

Thank god the CI tests finally passed!
@anirudh2290 A JIRA epic has already been created to track this project: https://issues.apache.org/jira/browse/MXNET-1184. I will bring this proposal to dev list for discusssion. Since this change should update APIs in all language bindings, help is needed from experts in different language support.