Skip to content
This repository has been archived by the owner on Nov 17, 2023. It is now read-only.

Optimize 'take' operator for CPU #20745

Merged
merged 7 commits into from
Jan 19, 2022
Merged

Optimize 'take' operator for CPU #20745

merged 7 commits into from
Jan 19, 2022

Conversation

bgawrych
Copy link
Contributor

Description

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

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
Copy link

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.

@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-work-in-progress PR is still work in progress and removed pr-awaiting-testing PR is reviewed and waiting CI build and test pr-work-in-progress PR is still work in progress labels Nov 19, 2021
@bgawrych
Copy link
Contributor Author

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

@mxnet-bot
Copy link

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

@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-work-in-progress PR is still work in progress and removed pr-work-in-progress PR is still work in progress pr-awaiting-testing PR is reviewed and waiting CI build and test labels Nov 23, 2021
@bgawrych
Copy link
Contributor Author

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

@mxnet-bot
Copy link

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

@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-work-in-progress PR is still work in progress and removed pr-work-in-progress PR is still work in progress pr-awaiting-testing PR is reviewed and waiting CI build and test labels Nov 23, 2021
@bgawrych
Copy link
Contributor Author

@mxnet-bot run ci [unix-gpu]

@mxnet-bot
Copy link

Jenkins CI successfully triggered : [unix-gpu]

@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-work-in-progress PR is still work in progress and removed pr-work-in-progress PR is still work in progress pr-awaiting-testing PR is reviewed and waiting CI build and test labels Nov 24, 2021
@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-awaiting-review PR is waiting for code review and removed pr-work-in-progress PR is still work in progress pr-awaiting-testing PR is reviewed and waiting CI build and test labels Dec 16, 2021
@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-awaiting-review PR is waiting for code review and removed pr-awaiting-review PR is waiting for code review pr-awaiting-testing PR is reviewed and waiting CI build and test labels Jan 4, 2022
Co-authored-by: bartekkuncer <bartosz.kuncer@intel.com>
@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-work-in-progress PR is still work in progress and removed pr-awaiting-review PR is waiting for code review pr-awaiting-testing PR is reviewed and waiting CI build and test labels Jan 18, 2022
@bgawrych
Copy link
Contributor Author

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

@mxnet-bot
Copy link

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

@mseth10 mseth10 added pr-awaiting-testing PR is reviewed and waiting CI build and test pr-awaiting-review PR is waiting for code review and removed pr-work-in-progress PR is still work in progress pr-awaiting-testing PR is reviewed and waiting CI build and test labels Jan 18, 2022
@bgawrych bgawrych merged commit 69e6c04 into apache:master Jan 19, 2022
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
pr-awaiting-review PR is waiting for code review
Projects
None yet
Development

Successfully merging this pull request may close these issues.

7 participants