torch.eye(d) is slow and hogs cpu for d >= 182

[akamaus]

🐛 Bug

I stumbled upon excessive CPU usage for my training code running on GPU. After some investigations I found the culprit.
It basically was

x = torch.eye(256).to('cuda') 

To Reproduce

This is quick and loads single CPU core.

%%timeit
    torch.eye(181)
6.43 µs ± 218 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

This is 3 times slower and hogs 30 cores.

%%timeit
    torch.eye(182)
21.2 µs ± 84.8 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Environment

I tested on server with 72 Intel(R) Xeon(R) Gold 6140 CPU @ 2.30GHz cores (hyper-threading is on)

% python3 collect_env.py

Collecting environment information...
PyTorch version: 1.6.0
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.3 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2

Python version: 3.6 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration:
GPU 0: GeForce RTX 2080 Ti
GPU 1: GeForce RTX 2080 Ti
GPU 2: GeForce RTX 2080 Ti
GPU 3: GeForce RTX 2080 Ti
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti
GPU 7: GeForce RTX 2080 Ti

Nvidia driver version: 440.100
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.18.2
[pip3] pytorch-wavelets==1.2.2
[pip3] torch==1.6.0
[pip3] torchvision==0.4.2
[conda] Could not collect

cc @VitalyFedyunin @ngimel

[ngimel]

To workaround, can you try torch.set_num_threads(1) ? We should look at the thread settings in eye, something goes wrong as multi-threading should not lead to a slowdown.

[akamaus]

@ngimel
As a workaround I'm perfectly happy with torch.eye(256, device='cuda'), and main problem is not a slowdown but heavy CPU load, so I can't run several experiments simultaneously. By the way. Even now I see some rare sparks of CPU load (%cpu jumps from 100 to 2000 in top). Is there a way to find the culprits?

[imaginary-person]

@ngimel
As a workaround I'm perfectly happy with torch.eye(256, device='cuda'), and main problem is not a slowdown but heavy CPU load, so I can't run several experiments simultaneously. By the way. Even now I see some rare sparks of CPU load (%cpu jumps from 100 to 2000 in top). Is there a way to find the culprits?

The OpenMP implementation uses a threshold at::internal::GRAIN_SIZE to determine whether parallelization would be helpful for an array, or not, and it is deemed inefficient to parallelize over arrays shorter than GRAIN_SIZE. GRAIN_SIZE has heuristically been chosen as 32768 in the current implementation.

I had initially assumed that d >= 182 triggers more threads because 182 * 182 is 33124, which being greater than 32768, triggers OpenMP parallelism, and d = 181 doesn't. However, my assumption was incorrect.
Since an identity matrix can be created by simply filling d number of 1s in a zeroed-matrix, only d elements need to be assigned a value of 1. So, for this operation, work has to be done on only d number of elements, which is what at:parallel_for() does when called from eye_out_cpu().

But neither d = 181 nor d = 182 should result in parallelizing work across multiple threads as both fail OpenMP's parallelization criterion (#pragma omp parallel if), as they're less than GRAIN_SIZE.

So, I'm not sure why d = 182 results in more threads being used than d = 181. I don't know what the additional threads do, but I observed this behavior with strace as well.
I'll try to debug with a debug build that I can use to set some breakpoints. It'd be great if someone could give me some advice on how I can debug this issue. Thanks!

[imaginary-person]

The OpenMP implementation uses a threshold at::internal::GRAIN_SIZE to determine whether parallelization would be helpful for an array, or not, and it is deemed inefficient to parallelize over arrays shorter than GRAIN_SIZE. GRAIN_SIZE has heuristically been chosen as 32768 in the current implementation.

I had initially assumed that d >= 182 triggers more threads because 182 * 182 is 33124, which being greater than 32768, triggers OpenMP parallelism, and d = 181 doesn't. However, my assumption was incorrect.

I'll try to debug with a debug build that I can use to set some breakpoints. It'd be great if someone could give me some advice on how I can debug this issue. Thanks!

I checked with a debug build.

OBSERVATIONS

1. d >= 182 results in OpenMP creating additional threads for zeroing 182 * 182 elements of the matrix using the at::native::zero_() method, as 33124 exceeds GRAIN_SIZE, but 181 * 181 does not. The additional threads correspond to OpenMP threads for zeroing the matrix. For d less than 182, elements are zeroed via std::memset().
2. Although the number of threads created by OpenMP are equal to the number of physical cores, only two of them end up getting any work assigned in the parallel for loop if d = 182.
3. As a sidenote, the at::parallel_for() in the eye operation itself doesn't utilize OpenMP parallelism, as it has to write to just 182 memory locations, which is less than GRAIN_SIZE.

POTENTIAL SOLUTION:
A higher threshold than at::internal::GRAIN_SIZE can probably be used for the zero_() method, if overhead due to parallelism actually hurts performance in this case, which would be more pronounced on machines with a large number of cores, and/or multiple NUMA nodes, as OpenMP creates threads equal to the number of the cores by default:

The following threshold can be adjusted:

pytorch/aten/src/ATen/native/Fill.cpp

Line 110 in 0ea1abe

nelements < internal::GRAIN_SIZE) {

.

However, choosing a threshold would again be based more on heuristics yet again, as no one value would suit all machines.
Perhaps the decision to set a higher threshold can take into account:

1. number of elements - already being used.
2. dtype().itemsize()- the size of the types of the elements
   1 and 2 should be multiplied, as they're a proxy for the number of cache lines that would be loaded into the processor cache(s), in order to be zeroed.

@VitalyFedyunin : Please advise if you already have a fix in mind for this issue, and if I'm missing something. Thanks!

[sparkingdark]

is this resolved?

[Bhavay-2001]

@janeyx99 , is this issue resolved??
Can I work on this issue. Please suggest some resoruces to resolve for the same. Thanks

[janeyx99]

@janeyx99 , is this issue resolved?? Can I work on this issue. Please suggest some resoruces to resolve for the same. Thanks

Hey @Bhavay192! I don't think this issue has been resolved as @imaginary-person had closed the linked attempt (not merged). However, I do not have the expertise to propose resources to solve this particular problem. Would @ngimel @VitalyFedyunin @imaginary-person and maybe @peterbell10 make better suggestions?

[aorenste]

Looked into this and adding some notes: I think there's a secondary problem which is that the iterator interface is very slow. Forcing a single thread (so no thread startup/queuing overhead) is 60x slower than the direct call to at::native::zero_().