Speed-up "advanced" indexing operations (#13420)
authorSam Gross <sgross@fb.com>
Tue, 27 Nov 2018 23:18:39 +0000 (15:18 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Tue, 27 Nov 2018 23:23:59 +0000 (15:23 -0800)
commit006505bb8f9dcf0f38b32518308d071c8a1ccec6
tree0abac8f0f7e2e4b7eeb161246267bb6c24a69c8e
parent0199d59d3a3ba69c62480ddb6680c9c5f55667b6
Speed-up "advanced" indexing operations (#13420)

Summary:
This speeds-up "advanced" indexing (indexing a tensor by a tensor)
on CPU and GPU. There's still a bunch of work to do, including
speeding up indexing by a byte (boolean) mask and speeding up the derivative
calculation for advanced indexing.

Here's some speed comparisons to indexing on master using a little [benchmark script](https://gist.github.com/colesbury/c369db72aad594e5e032c8fda557d909) with 16 OpenMP threads and on a P100. The test cases are listed as (input shape -> output shape).

| Test case             | CPU (old vs. new)   | CUDA (old vs. new)     |
|-----------------------|---------------------|------------------------|
| 1024x1024 -> 512x1024 | 225 us vs. **57 us**  | 297 us vs. **47 us** |
| 1024x1024 -> 1024x512 | 208 us vs. **153 us** | 335 us vs. **54 us** |
| 50x50 -> 20000x50     | 617 us vs. **77 us**  | 239 us vs. **54 us** |
| 50x50 -> 50x20000     | 575 us vs. **236 us** | 262 us vs. **58 us** |
| 2x5x10 -> 10          | 65 us  vs. **18 us**  | 612 us vs. **93 us** |

See #11647
Pull Request resolved: https://github.com/pytorch/pytorch/pull/13420

Reviewed By: soumith

Differential Revision: D13088936

Pulled By: colesbury

fbshipit-source-id: 0a5c2ee9aa54e15f96d06692d1694c3b24b924e2
21 files changed:
aten/src/ATen/core/Tensor.h
aten/src/ATen/core/TensorMethods.h
aten/src/ATen/core/Type.h
aten/src/ATen/cuda/detail/OffsetCalculator.cuh
aten/src/ATen/native/Indexing.cpp
aten/src/ATen/native/Indexing.h [new file with mode: 0644]
aten/src/ATen/native/TensorIterator.cpp
aten/src/ATen/native/TensorIterator.h
aten/src/ATen/native/cpu/IndexKernel.cpp [new file with mode: 0644]
aten/src/ATen/native/cuda/IndexKernel.cu [new file with mode: 0644]
aten/src/ATen/native/native_functions.yaml
test/expect/TestScript.test_index_put_trace_with_view.expect
test/expect/TestScript.test_index_put_trace_without_view.expect
test/run_test.py
test/test_cuda.py
test/test_indexing.py
test/test_indexing_cuda.py [new file with mode: 0644]
tools/autograd/derivatives.yaml
torch/csrc/autograd/VariableTypeManual.cpp
torch/csrc/autograd/VariableTypeUtils.h
torch/onnx/symbolic.py