[FFI] Part1: npx.softmax, npx.activation, npx.batch_norm, npx.fully_connected

[barry-jin]

Description

Adopt packedfunc based ffi on some frequently used numpy_extension operators. Some benchmarks are as follows:
npx.softmax and npx.log_softmax

setup = """
from mxnet import np, npx
npx.set_np()
a = np.ones((2, 2))
"""
stmt = """
npx.softmax(a)
"""
legacy
7.195554300000006e-05
New FFI
5.542412399999996e-05


setup = """
from mxnet import np, npx
npx.set_np()
a = np.array([1, 2, .1])
"""
stmt = """
npx.log_softmax(a)
"""
legacy
7.109741400000047e-05
New FFI
5.514604400000067e-05

npx.activation

setup = """
from mxnet import np, npx
npx.set_np()
x = np.arange(8).reshape((2, 2, 2))
"""
stmt = """
o = npx.activation(x, act_type='sigmoid')
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00010351458600000019
New FFI
4.998078599999989e-05

npx.batch_norm

setup = """
from mxnet import np, npx
npx.set_np()
a = np.zeros((2,2))
b = np.ones((2,2))
c = np.ones((2,2))
d = np.ones((2,2))
e = np.ones((2,2))
"""
stmt = """
o = npx.batch_norm(a, b, c, d, e,
                   eps=1e-3,
                   momentum=0.9,
                   fix_gamma=True,
                   use_global_stats=False,
                   output_mean_var=False,
                   axis=1,
                   cudnn_off=False,
                   min_calib_range=None,
                   max_calib_range=None)
"""

timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00017782039200000012
New FFI
8.22070250000002e-05

npx.fully_connected

setup = """
from mxnet import np, npx
npx.set_np()
data = np.arange(8).reshape((4,2))
weight = np.arange(4).reshape((2,2))
bias = np.arange(2).reshape((2,))
num_hidden = bias.shape[0]
"""
stmt = """
o = npx.fully_connected(data, weight, bias, num_hidden=num_hidden, no_bias=False)
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
9.581770600000006e-05
New FFI
6.12347289999997e-05

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 @barry-jin , 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: [centos-gpu, unix-gpu, miscellaneous, website, sanity, edge, windows-gpu, unix-cpu, windows-cpu, clang, 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.

[barry-jin]

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

[mxnet-bot]

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

[barry-jin]

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

[mxnet-bot]

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

[barry-jin]

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

[mxnet-bot]

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