[FFI] part2: npx.pick, npx.convolution, npx.deconvolution

[barry-jin]

Description

part2 of #20096
Benchmarks:
npx.pick

setup = """
from mxnet import np, npx
npx.set_np()
pred = np.arange(6).reshape((3,2))
label = np.array([2, -1, -2])
"""
stmt = """
o = npx.pick(pred, label, axis=1, keepdims=True, mode='wrap')"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.0001114548919999998
New FFI
7.543004800000031e-05

npx.convolution

setup = """
from mxnet import np, npx
npx.set_np()
pred = np.arange(6).reshape((3,2))
dim = 2
batch_size = 1
channel = 3
height = 5
width = 9 // 3
num_filter = 4
kernel = (3,) * dim   # => shape = (3, 3)
inp = np.ones(shape=(batch_size, channel, height, width))
weight = np.ones(shape=(num_filter, channel, kernel[0], kernel[1]))
bias = np.array(num_filter,)
"""
stmt = """
out = npx.convolution(data=inp, weight=weight, num_filter=num_filter, \
                      kernel=kernel, stride=(height, height), no_bias=True, \
                      layout='NCHW', num_group=1, \
                      workspace=1024, cudnn_tune="off", cudnn_off=True)
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.0001754432060000002
New FFI
9.782946400000014e-05

npx.deconvolution

setup = """
from mxnet import np, npx
npx.set_np()
pred = np.arange(6).reshape((3,2))
dim = 2
batch_size = 1
channel = 3
height = 5
width = 9 // 3
num_filter = 4
kernel = (3,) * dim   # => shape = (3, 3)
inp = np.ones(shape=(batch_size, channel, 1, width))
weight = np.ones(shape=(channel, num_filter, kernel[0], kernel[1]))
bias = np.array(num_filter,)
"""
stmt = """
out = npx.deconvolution(data=inp, weight=weight, num_filter=num_filter, \
                      kernel=kernel, stride=(height, height), no_bias=True, \
                      layout='NCHW', num_group=1, \
                      workspace=1024, cudnn_tune="off", cudnn_off=True)
"""
timer = timeit.Timer(setup=setup,
                     stmt=stmt)
num_repeat = 1000
print(min(timer.repeat(repeat=10, number=num_repeat)) / num_repeat)

legacy
0.00016225327699999958
New FFI
0.00010197798500000043

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

---

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.