Description
🐛 Bug
Reduce operations that depend on binary_kernel_reduce suffers from huge slowdowns (Upto 100x in most cases) whenever the dimension over which these operations are applied is small compared to the other dimensions. These include torch.norm / torch.linalg.norm, torch.std, torch.var, torch.argmax, torch.argmin.
Apart from the above ones, torch.mean, torch.sum, torch.max, torch.min suffer from less severe performance losses. The code snippet below illustrates the performance gaps:
In [34]: inp = torch.rand(10**6, 2) # Similar observations for higher dimensions.
In [35]: %timeit torch.sum(inp)
210 µs ± 2.86 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [36]: %timeit torch.sum(inp, dim=0)
998 µs ± 28.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [37]: %timeit torch.sum(inp, dim=1)
4.78 ms ± 61.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [43]: %timeit torch.max(inp)
306 µs ± 932 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [44]: %timeit torch.max(inp, dim=0)
1.58 ms ± 40.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [45]: %timeit torch.max(inp, dim=1)
4.08 ms ± 6.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [47]: %timeit torch.norm(inp)
1.46 ms ± 15.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [47]: %timeit torch.norm(inp, dim=0)
1.59 ms ± 9.18 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [48]: %timeit torch.norm(inp, dim=1)
146 ms ± 415 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [50]: %timeit torch.mean(inp)
248 µs ± 1.74 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [51]: %timeit torch.mean(inp, dim=0)
1.07 ms ± 7.57 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [52] %timeit torch.mean(inp, dim=1)
4.87 ms ± 25.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
I have been studying the implementations and been looking for the bottlenecks. It turns out that the TensorIteratorBase::foreach_reduced_elt
pytorch/aten/src/ATen/native/TensorIteratorReduce.cpp
Lines 121 to 172 in 93973ee
binary_kernel_reduce relies upon is the bottleneck which is mainly due to the following while loop : pytorch/aten/src/ATen/native/TensorIteratorReduce.cpp
Lines 143 to 148 in 93973ee
inp = torch.rand(10**6, 2)
torch.norm(inp, dim=1)If N threads are spawned, then each thread iterates (10**6) / N times in the above loop which is not great given that the operations performed inside the loop are not lightweight! I have done some profiling here. There should be a faster approach which is rather raw instead of using DimCounter within TensorIteratorBase::foreach_reduced_elt. That method will likely need redesigning. Similar observations were found regarding torch.norm. I will be happy to work on this and can send a PR.
Environment
PyTorch version: 1.8.0a0+508bab4
Is debug build: False
CUDA used to build PyTorch: Only CPU
ROCM used to build PyTorch: N/A
OS: Linux Mint 20 (x86_64)
GCC version: (Ubuntu 10.2.0-5ubuntu1~20.04) 10.2.0
Clang version: Could not collect
CMake version: version 3.16.3
Python version: 3.7 (64-bit runtime)
HIP runtime version: N/A
MIOpen runtime version: N/A
Versions of relevant libraries:
[pip3] numpy==1.21.0.dev0+430.ge745a19cb
[pip3] torch==1.8.0a0+unknown
[conda] Could not collect
cc @VitalyFedyunin @ngimel