Generalize catArray for contiguous inputs and dim != 0 (#17032)
authorJames Reed <jamesreed@fb.com>
Thu, 14 Feb 2019 23:58:06 +0000 (15:58 -0800)
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>
Fri, 15 Feb 2019 00:33:23 +0000 (16:33 -0800)
commitf1da9892e94acc8578bc17567e9618f5fbf6cb8d
treed1de67d47d3ffb3323e24dc5e50031d4c53750ab
parentf3dd5563e44616bc2be5317cdfd7947849e2efd8
Generalize catArray for contiguous inputs and dim != 0 (#17032)

Summary:
I noticed that we were sinking a lot of time into `cat` operations in machine translation on CPU, and drilled down to us doing the cat element-by-element, even though all the inputs were contiguous. The reason was we were doing the cat along a dimension that was not 0, and that caused us to not use the fast `memcpy` branch. This PR generalizes that branch.

Quick benchmark script:
```
import torch, time

tensors = [torch.rand(6, 2, 1024) for i in range(5)]

NITER = 1000
s = time.time()
for i in range(NITER):
    torch.cat(tensors, dim=1)
print('time per iter ', (time.time() - s) / NITER)
```

Before:
```
time per iter  8.089399337768554e-05
```

After:
```
time per iter  2.183413505554199e-05
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17032

Differential Revision: D14090038

Pulled By: jamesr66a

fbshipit-source-id: 2c733a84915896008ac95f2233f44894bd2573de
aten/src/TH/generic/THTensor.cpp