Add ATen pdist CPU kernel by erikbrinkman · Pull Request #10782 · pytorch/pytorch

erikbrinkman · 2018-08-22T18:34:49Z

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

torch/nn/functional.py

test/test_jit.py

ezyang · 2018-08-22T18:44:05Z

How does the speed compare to scipy? ;)

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

ezyang · 2018-08-22T18:58:24Z

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

erikbrinkman · 2018-08-22T19:04:34Z

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 · 2018-08-22T20:18:13Z

@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).

torch/nn/functional.py

test/test_nn.py

test/test_torch.py

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

torch/nn/functional.py

test/test_torch.py

ezyang

But don't land until tests pass...

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

Notably this now enforces a constract that the input is contiguous.

facebook-github-bot

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

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

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

erikbrinkman requested review from apaszke, colesbury, ezyang, gchanan, soumith and zdevito as code owners August 22, 2018 18:34

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

gchanan reviewed Aug 22, 2018

View reviewed changes

ssnl reviewed Aug 24, 2018

View reviewed changes

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 approved these changes Aug 28, 2018

View reviewed changes

erikbrinkman force-pushed the master branch 2 times, most recently from c8142db to 4b747e8 Compare August 29, 2018 07:35

erikbrinkman added 8 commits August 30, 2018 00:40

Add ATen pdist CPU kernel

94c1e6a

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

Address some of the PR comments and fix some python lint

0c1c2d1

Address more PR issues

50752d3

Update pdist tests to make 0 and inf norm more clear

896e82f

Refactor pdist call surface to better support CUDA kernel

96c42a4

Notably this now enforces a constract that the input is contiguous.

Refactor kernel again on inlining recommendations

c4187d4

Attempt to debug autograd errors on other environments

45a6477

Fix bug when parallel spans more than one column

0d26aa6

erikbrinkman force-pushed the master branch from e0b9cda to 0d26aa6 Compare August 30, 2018 07:41

facebook-github-bot reviewed Aug 30, 2018

View reviewed changes

facebook-github-bot closed this in 611a608 Aug 30, 2018

vadimkantorov mentioned this pull request Aug 30, 2018

[pytorch] [feature request] Pairwise distances between all points in a set (a true pdist) #9406

Closed

ezyang added the merged label Jun 26, 2019

Conversation

erikbrinkman commented Aug 22, 2018

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

ezyang commented Aug 22, 2018

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

ezyang commented Aug 22, 2018

Uh oh!

erikbrinkman commented Aug 22, 2018

Uh oh!

cpuhrsch commented Aug 22, 2018 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

This comment was marked as off-topic.

Uh oh!

cpuhrsch commented Aug 22, 2018 •

edited

Loading