OpenBLAS OpenMP support AArch64 builds.

[nSircombe]

Currently, AArch64 builds rely on a single-threaded build of OpenBLAS, see:

builder/build_aarch64_wheel.py

Line 182 in 8e799eb

host.run_cmd("pushd OpenBLAS; make NO_SHARED=1 -j8; sudo make NO_SHARED=1 install; popd")

Inclusion of OpenMP is marked as a TODO.

Enabling support should be a matter of adding the USE_OPENMP=1 flat at

builder/build_aarch64_wheel.py

Line 182 in 8e799eb

host.run_cmd("pushd OpenBLAS; make NO_SHARED=1 -j8; sudo make NO_SHARED=1 install; popd")

.

Note: When building PyTorch with an OpenBLAS myself, I also explicitly set OpenBLAS_HOME='/opt/OpenBLAS' BLAS='OpenBLAS' USE_MKLDNN=0 USE_OPENMP=1 USE_LAPACK=1 when running setup.py.

Can OpenMP support be enabled in build_aarch64_wheel.py?
Are there any issues currently blocking this change?

Note: this mirrors an issue raised in #679, however it is unrelated to the issues concerning choices of mcpu, mtune, and march, so I felt it would be beneficial to separate it out and address any specific issues separately.

[malfet]

@nSircombe last time I've tried enabling OpenMP support in OpenBLAS and linking statically, I run into can not allocate LTS runtime issue, but let me try again.

[nSircombe]

Hi @malfet,
That's interesting, I have not hit that issue myself - I'd be very interested to know how you get on. If we can reproduce then I'm happy to dig a little deeper.

[malfet]

Hi @nSircombe,
I was about to write, that I can not reproduce the problem, but the reason was that pretty simple: when I've enabled OpenMP in OpenBLAS, pytorch started to use Eigen instead. I use the following litmus test to check whether compilation with OpenBLAS was successful or not:

$ python -c "import torch;x=torch.rand(3,3);print(torch.__version__, torch.svd(torch.mm(x,x.t())))"
1.9.0.dev20210503 torch.return_types.svd(
U=tensor([[-0.5307, -0.0179, -0.8474],
        [-0.6476, -0.6365,  0.4190],
        [-0.5469,  0.7711,  0.3262]]),
S=tensor([1.7903, 0.2885, 0.0034]),
V=tensor([[-0.5307, -0.0179, -0.8474],
        [-0.6476, -0.6365,  0.4190],
        [-0.5469,  0.7711,  0.3262]]))

By the way, can you share a code snippet that triggers "Detect OpenMP Loop and this application may hang." warnings?

[nSircombe]

HI @malfet,

Your 'smoketest' works for me on my OpenMP-ed from-src build. I notice that the torch.__config__ on the released whl shows USE_EIGEN_FOR_BLAS=ON - this is not the case for my from-builds source.

can you share a code snippet that triggers "Detect OpenMP Loop and this application may hang."

Yep, I'll tidy it up and add something here.

[nSircombe]

Here's a reproducer for the "Detect OpenMP Loop and this application may hang." warning.

import torch
import urllib
import torchvision.models as models
from PIL import Image
from torchvision import transforms
 
url, fname = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
urllib.request.urlretrieve(url, fname)
input_img = Image.open(fname)
pp_img = transforms.Compose([
  transforms.Resize(256),
  transforms.CenterCrop(224),
  transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406],
  std=[0.229, 0.224, 0.225]),
])

input_ten = pp_img(input_img)
input_bat = input_ten.unsqueeze(0)

model = models.resnet50(pretrained=True)
model.eval()

out = model(input_bat)

Running with anything other than OMP_NUM_THREADS=1 with the 1.8.1 wheel produces the warning for me. Builds with an OpenMP enabled OpenBLAS do not.

[AGSaidi]

@nSircombe what compiler is openblas being built with. Is it being built with -moutline-atomics? We certainly saw issues with OpenBLAS pthreads with ld/st exclusives staving threads.

[nSircombe]

I've built with GCC 9.3 and 10.2, I'm not using -moutline-atomics.

[malfet]

Looks like TLS problem was fixed in OpenBLAS 0.3.15, so will try to enable the compilation OpenMP enabled OpenBLAS for 1.9.0 see pytorch/pytorch#59428 and appropriate builder change

[nSircombe]

Hi @malfet,
How did you get on with the +OpenMP build following this change?
Does the resulting build pickup OpenBLAS as expected (and does not show USE_EIGEN_FOR_BLAS=ON in the config?)

[nSircombe]

and appropriate builder change

...would you be able to point me towards the builder change?

[malfet]

@nSircombe I've already made changes to builder (by just adding USE_OPENMP) and 1.9.2-rc2 builds should be available from https://download.pytorch.org/whl/test/cpu/torch_test.html

Would appreciate if you can smoke test them and let me know if they show perf gains you've expected from enabling OpenMP.

[nSircombe]

Great, thanks @malfet, I'll take the new whl for a spin and let you know how I get on.

[nSircombe]

Hi @malfet,

I've installed the whl from that url (pip install torch -f https://download.pytorch.org/whl/test/cpu/torch_test.html) but it looks like the whl has still picked up Eigen. torch.__config__.show() lists USE_EIGEN_FOR_BLAS=ON, and a simple check does not show torch using all the available threads (https://forums.developer.nvidia.com/t/pytorch-and-numpy-only-run-with-one-core/110430) - in my experience this happens when the linkage to OpenBLAS has quietly failed.

For my own builds from src. I'm currently using a build line like: OpenBLAS_HOME=$OPENBLAS_DIR BLAS="OpenBLAS" USE_OPENMP=1 USE_LAPACK=1 USE_CUDA=0 USE_FBGEMM=0 USE_DISTRIBUTED=0 python setup.py install and the result does not have USE_EIGEN_FOR_BLAS=ON (and with BLAS="OpenBLAS" set explicitly, fails if unable to link to OpenBLAS)

[malfet]

Hi @nSircombe,

Not sure how you've measured CPU utilization in the example you've referred to, I rely on time tool, which seem to show that torch.bmm is indeed parallelizable across cores:

$ uname -m; time python3 -c "import torch; print(torch.__version__, torch.get_num_threads());a=torch.rand(100, 100, 100); b=torch.rand(100,100, 100); [torch.bmm(a,b).sum() for i in range(1000)]"
aarch64
1.9.0 8

real	0m3.323s
user	0m23.233s
sys	0m0.168s

user time is roughly 8x larger than real time(and torch.get_num_threads() returned 8), which means torch.bmm was running on all available cores. And USE_EIGEN_FOR_BLAS=ON in the config is a red herring, if it were used, torch.svd, which relies on LAPACK code from OpenBLAS would not have worked.

[nSircombe]

Hi @malfet

torch.svd, which relies on LAPACK code from OpenBLAS would not have worked.

...that's a good point.

Not sure how you've measured CPU utilization...

I was just eye-balling htop while running some simple tests.
When I look at the output from your simple example, I see it reporting the current number of threads, but I only see 8 active in htop, and the timing is consistent with 8 threads. Even if I explicitly torch.set_num_threads(32) for example:

> uname -m; time python3 -c "import torch; torch.set_num_threads(32); print(torch.__version__, torch.get_num_threads());a=torch.rand(100, 100, 100); b=torch.rand(100,100, 100); [torch.bmm(a,b).sum() for i in range(1000)]"     
aarch64
1.9.0 32

real    0m4.542s
user    0m34.504s
sys     0m3.137s