Package libomp.so with torchaudio wheel for ROCm

[jithunnair-amd]

Problem:
#2485 introduced a build-time dependency on libomp.so for torchaudio, but when a torchaudio wheel built with that PR is imported in an environment where ROCm is NOT installed, it errors out e.g.:
https://github.com/pytorch/builder/actions/runs/5954772894/job/16152043224

  File "/opt/conda/envs/conda-env-5954772894/lib/python3.11/site-packages/torchaudio/__init__.py", line 1, in <module>
    from . import (  # noqa: F401
  File "/opt/conda/envs/conda-env-5954772894/lib/python3.11/site-packages/torchaudio/_extension/__init__.py", line 45, in <module>
    _load_lib("libtorchaudio")
  File "/opt/conda/envs/conda-env-5954772894/lib/python3.11/site-packages/torchaudio/_extension/utils.py", line 64, in _load_lib
    torch.ops.load_library(path)
  File "/opt/conda/envs/conda-env-5954772894/lib/python3.11/site-packages/torch/_ops.py", line 839, in load_library
    ctypes.CDLL(path)
  File "/opt/conda/envs/conda-env-5954772894/lib/python3.11/ctypes/__init__.py", line 376, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: libomp.so: cannot open shared object file: No such file or directory

Proposed solution:

PART 1:
We need to package the libomp.so library (that comes as part of ROCm installation) with the torchaudio wheel (similar to how we package ROCm libraries with torch wheel). This PR adds the libomp.so to the wheel packaging.

PART 2:
The above solution will be needed in conjunction with another fix in the wheel building yml file (this or this?) to set the RUNPATH for libtorchaudio.so so that it looks in its current directory ($ORIGIN) for any dependencies. This is a technique already used in the PyTorch wheels here.

However, I propose just adding $ORIGIN to the RUNPATH instead of replacing it, like we do in the PyTorch wheels (for minimal behavior change):

export PATCHELF_BIN=/usr/local/bin/patchelf (already available in the manylinux-builder docker image)
CURRENT_RPATH=$($PATCHELF_BIN --print-rpath libtorchaudio.so)
$PATCHELF_BIN --set-rpath '$ORIGIN:'$CURRENT_RPATH libtorchaudio.so

This brings the following change in RUNPATH for libtorchaudio.so:

<< /opt/rocm/llvm/lib:/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/lib:/opt/rocm/hip/lib:/opt/rocm/roctracer/lib:/opt/rocm/lib
====
>> $ORIGIN:/opt/rocm/llvm/lib:/opt/conda/envs/py_3.8/lib/python3.8/site-packages/torch/lib:/opt/rocm/hip/lib:/opt/rocm/roctracer/lib:/opt/rocm/lib

ldd shows the expected output:

[root@9d581fa3e883 lib]# ldd libtorchaudio.so
./libtorchaudio.so: /lib64/libm.so.6: version `GLIBC_2.27' not found (required by ./libtorchaudio.so)
        libomp.so => /opt/_internal/cpython-3.8.1/lib/python3.8/site-packages/torchaudio/lib/./libomp.so (0x00007fd8f40b7000)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/audio/3575

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 11 Pending

As of commit e9f5c88 with merge base 6fb6854 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jithunnair-amd]

cc @atalman The above is an attempt to fix the binary import test failure for torchaudio

[jithunnair-amd]

@pruthvistony @mthrok

[jithunnair-amd]

@atalman This is part 1 of my proposed solution for fixing the torchaudio wheel issue with libomp.so. However, it will also need a part 2 fix in the yml file for building the wheel. Can you please take a look at the proposed solution in the PR description and let us know if it sounds okay?

[atalman]

this will affect all builds not only rocm. Since this issue happens only on rocm can we implement this only for rocm builds ?

[jithunnair-amd]

Actually this should be torch.version.hip is not None. That should ensure that this change is applied only to ROCm builds. Doest that sound good?
@pruthvistony, can you please make the changes as appropriate?

[pruthvistony]

Changes pushed. It should affect only ROCm builds now.

[mthrok]

Is there a way to handle this in PyTorch core in abstracted manner, like how CUDA is handled?

[jithunnair-amd]

Is there a way to handle this in PyTorch core in abstracted manner, like how CUDA is handled?

Since libomp.so is not a dependency for Pytorch, I didn't consider adding it to Pytorch wheels. I'll need @atalman to weigh in here if that's the route we should go. It'll probably remove any need for changes to torchaudio wheel recipe, since the torchaudio so should be able to pick up libomp.so from torch wheel install path in that case.

@sunway513 @pruthvistony

[atalman]

What is the test plan for this ?

[jithunnair-amd]

This is a fix for the import of torchaudio using wheel, so in a way, we already have a test case for this, right?

[atalman]

@pruthvistony @jithunnair-amd for second part of the fix we should do it in torchaudio as well, via post script. Since this is audio only change we should not put it in test-infra : https://github.com/pytorch/test-infra/blob/65abc6ecc3afe43012e74bc45b824f6cd36c1eff/.github/workflows/build_wheels_linux.yml#L187

[mthrok]

Is there a way to handle this in PyTorch core in abstracted manner, like how CUDA is handled?

Since libomp.so is not a dependency for Pytorch, I didn't consider adding it to Pytorch wheels. I'll need @atalman to weigh in here if that's the route we should go. It'll probably remove any need for changes to torchaudio wheel recipe, since the torchaudio so should be able to pick up libomp.so from torch wheel install path in that case.

@sunway513 @pruthvistony

OpenMP is used by many different projects, so I'd be careful with shipping libomp.so. In the past, we had an issue where libsox uses OMP which interferes with MKL OMP used by PyTorch. I think at least we need to find a way so that PyTorch and torchaudio agree on what libomp they use. Packaging and shipping libomp in torchaudio does not sound like a comprehensive solution to me.