Optimize 'take' operator for CPU

[bgawrych]

Description

This change optimizes take operator by new CPU Kernel - it is performing less memory operations (IO), especially for axes in the middle.

Benchmark on CLX-8280 (28 cores)
Script:

import mxnet
import mxnet.gluon.nn as nn
import mxnet.numpy as np
import time

dims = [128, 512, 1024, 4096]
print("shape;axis;time")
for ndim in range (2):
   for dim1 in dims:
     for dim2 in dims:
        shape = (dim1, dim2) if ndim == 0 else (32, dim1, dim2)
        a = np.random.uniform(-1.0, 1.0, shape).astype(np.float32)
        for axis in range(2 + ndim):
            indices = np.random.uniform(0, shape[axis], (shape[axis]//2,)).astype(np.int32)
            tic = time.time()
            for i in range(200):
                out = np.take(a, indices, axis=axis)
                out.wait_to_read()
            toc = time.time()
            print(f"{shape};{axis};{toc-tic}")

Checklist

Essentials

• PR's title starts with a category (e.g. [BUGFIX], [MODEL], [TUTORIAL], [FEATURE], [DOC], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

Changes

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

Comments

• If this change is a backward incompatible change, why must this change be made.
• Interesting edge cases to note here

[mxnet-bot]

Hey @bgawrych , Thanks for submitting the PR
All tests are already queued to run once. If tests fail, you can trigger one or more tests again with the following commands:

• To trigger all jobs: @mxnet-bot run ci [all]
• To trigger specific jobs: @mxnet-bot run ci [job1, job2]

---

CI supported jobs: [website, clang, unix-cpu, miscellaneous, windows-cpu, sanity, windows-gpu, centos-gpu, unix-gpu, edge, centos-cpu]

---

Note:
Only following 3 categories can trigger CI :PR Author, MXNet Committer, Jenkins Admin.
All CI tests must pass before the PR can be merged.

[bgawrych]

@mxnet-bot run ci [windows-gpu, unix-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-gpu, windows-gpu]

[bgawrych]

@mxnet-bot run ci [windows-gpu, unix-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-gpu, windows-gpu]

[bgawrych]

@mxnet-bot run ci [unix-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [unix-gpu]

[bgawrych]

@mxnet-bot run ci [unix-cpu, windows-gpu]

[mxnet-bot]

Jenkins CI successfully triggered : [windows-gpu, unix-cpu]