Skip to content

Add ATen pdist CPU kernel #10782

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
erikbrinkman wants to merge 8 commits into from
Closed

Add ATen pdist CPU kernel #10782

erikbrinkman wants to merge 8 commits into from

Conversation

@erikbrinkman
Copy link
Contributor

@erikbrinkman erikbrinkman commented Aug 22, 2018

Also add single grad whitelist to the jit test

Test Plan:

python -m unittest test_torch.TestTorch.test_pdist_{empty,special,scipy}
test_nn.TestNN.test_pdist test_jit.TestJitGenerated.test_nn_pdist

ezyang
ezyang reviewed Aug 22, 2018
View reviewed changes
torch/nn/functional.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

ezyang
ezyang reviewed Aug 22, 2018
View reviewed changes
test/test_jit.py Outdated

This comment was marked as off-topic.

Sign in to view

@ezyang
Copy link
Contributor

ezyang commented Aug 22, 2018

How does the speed compare to scipy? ;)

ezyang
ezyang reviewed Aug 22, 2018
View reviewed changes
aten/src/ATen/native/cpu/DistanceOpsKernel.cpp Outdated

This comment was marked as off-topic.

Sign in to view

ezyang
ezyang reviewed Aug 22, 2018
View reviewed changes
aten/src/ATen/native/cpu/DistanceOpsKernel.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

@ezyang
Copy link
Contributor

ezyang commented Aug 22, 2018

This looks very nice, and nicely parallelized. I'd say the main blocker is docs.

@erikbrinkman
Copy link
Contributor Author

erikbrinkman commented Aug 22, 2018

For timing:

Note that for an n x m input, space complexity is n^2, and time complexity is m n^2.

shape scipy torch
10 x 10 11.8 µs ± 280 ns 6.65 µs ± 23.2 ns
10 x 1000 41.3 µs ± 900 ns 9.41 µs ± 25.3 ns
1000 x 10 3.25 ms ± 94.8 µs 311 µs ± 59.6 µs
1000 x 1000 241 ms ± 6.1 ms 35.7 ms ± 88.7 µs
10000 x 1000 36.1 s ± 631 ms 3.63 s ± 7.87 ms

@cpuhrsch
Copy link
Contributor

cpuhrsch commented Aug 22, 2018 •
edited
Loading

@erikbrinkman - Does you code benefit from using AVX2 or AVX? When you add a file to native/kernel it'll recompile with each of those capabilities (i.e. -mavx and -mavx2) and dispatch depending on the CPU capability. But I don't see you make explicit use of those instructions. Did you time this to see whether it helps? (You can use the environment variables defined in ATen/native/DispatchStub.cpp. Sometimes, it's worse using those extended instruction sets. Also see "[Note SSE-AVX transitions]" if you see a significant slowdown (it might help explain it).

gchanan
gchanan reviewed Aug 22, 2018
View reviewed changes
torch/nn/functional.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

test/test_nn.py Outdated

This comment was marked as off-topic.

Sign in to view

test/test_torch.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

aten/src/ATen/native/cpu/DistanceOpsKernel.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

ssnl
ssnl reviewed Aug 24, 2018
View reviewed changes
aten/src/ATen/native/cpu/DistanceOpsKernel.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

torch/nn/functional.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

ezyang
ezyang reviewed Aug 28, 2018
View reviewed changes
test/test_torch.py Outdated

This comment was marked as off-topic.

Sign in to view

ezyang
ezyang approved these changes Aug 28, 2018
View reviewed changes
Copy link
Contributor

@ezyang ezyang left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But don't land until tests pass...

@erikbrinkman erikbrinkman force-pushed the master branch 2 times, most recently from c8142db to 4b747e8 Compare August 29, 2018 07:35
facebook-github-bot
facebook-github-bot reviewed Aug 30, 2018
View reviewed changes
Copy link
Contributor

@facebook-github-bot facebook-github-bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

erikbrinkman has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

zdevito pushed a commit to zdevito/ATen that referenced this pull request Aug 30, 2018
@erikbrinkman @facebook-github-bot
Add ATen pdist CPU kernel (#10782)
a887bc9 
Summary:
Also add single grad whitelist to the jit test
Pull Request resolved: pytorch/pytorch#10782

Reviewed By: ezyang

Differential Revision: D9583378

Pulled By: erikbrinkman

fbshipit-source-id: 069e5ae68ea7f3524dec39cf1d5fe9cd53941944
@vadimkantorov vadimkantorov mentioned this pull request Aug 30, 2018
Closed
PenghuiCheng pushed a commit to PenghuiCheng/pytorch that referenced this pull request Sep 11, 2018
@erikbrinkman @PenghuiCheng
Add ATen pdist CPU kernel (pytorch#10782)
446a6f7 
Summary:
Also add single grad whitelist to the jit test
Pull Request resolved: pytorch#10782

Reviewed By: ezyang

Differential Revision: D9583378

Pulled By: erikbrinkman

fbshipit-source-id: 069e5ae68ea7f3524dec39cf1d5fe9cd53941944
@ezyang ezyang added the merged label Jun 26, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@cpuhrsch cpuhrsch cpuhrsch left review comments

@gchanan gchanan gchanan left review comments

@ssnl ssnl ssnl left review comments

@facebook-github-bot facebook-github-bot facebook-github-bot left review comments

@ezyang ezyang ezyang approved these changes

@apaszke apaszke Awaiting requested review from apaszke

@colesbury colesbury Awaiting requested review from colesbury

@soumith soumith Awaiting requested review from soumith

@zdevito zdevito Awaiting requested review from zdevito

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

6 participants

@erikbrinkman @ezyang @cpuhrsch @gchanan @ssnl @facebook-github-bot