No disco_cuda_extension after installing

[wongchieh]

Dear NVIDIA,

Thanks for open this utils, I came across an error when I use this.

I am trying to run makani and torch-harmonics is one of its part. But I always get an warning:

"""
torch_harmonics/convolution.py:515: UserWarning: couldn't find CUDA extension, falling back to slow PyTorch implementation
warn("couldn't find CUDA extension, falling back to slow PyTorch implementation")

"""

When I look at torch_harmonics/convolution.py:515, I think it failed with :
"""
from disco_helpers import preprocess_psi
import disco_cuda_extension
"""

This bothered me for a long time, hope you could have some advice on this. I just followed the way in this git repo to install the torch-harmonics.

Best,
Cheih

[azrael417]

Hello Cheih,

thanks for reaching out. So in order to build the CUDA extension, it is best to build the repo from source. We are currently looking into options for building a wheel but because our dependency tree is rather complex (depending on torch which in turn depends on CUDA versions), we might need to wait for the wheel-next standard to be defined. In the meantime, please compile TH from scratch. This is how it works:

1. git clone TH into some directory <th_directory>
2. set the following env variables:

you always need this:
export FORCE_CUDA_EXTENSION=1

set the arches you need to compile for, it will cut down compilation time by a lot
export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0 10.0+PTX"
Then run the build process:
cd <th_directory>; pip install --no-build-isolation .

You need to have a recent torch 2.X installed. Please let me know if that works.

Best regards
Thorsten

[wongchieh]

Dear Thorsten,

Thanks for your reply. I've tried the seting 'export FORCE_CUDA_EXTENSION=1; export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0 10.0+PTX"', and then install it. While It also failed when running makani:
"""
from disco_helpers import preprocess_psi
import disco_cuda_extension
"""

It is hard to implement this fast version of TH. Hope the wheel that you talk about will be released soon, so that we could make the best use of it.

Best,
Chieh

[azrael417]

Ok, so lets debug this: when you install torch harmonics the way I described above and then, without installing makani, just do that above:

from disco_helpers import optimized_kernels_is_available
print(optimized_kernels_is_available())

should evaluate to True. Does that work? If yes, then when you install makani, it will cease to work, I think one (or both) of the following happens

• makani installs its own copy of torch, so that the pytorch used by makani is not the one you compiled torch harmonics against
• makani installs its own copy of torch harmonics, not the one you compiled manually.

In order to check what is going on, I would test if you can download and compile TH in isolation of makani. The compilation can take some time (20 minutes or so, depending on the hardware), if the install process is fast something did not work.
Then, once that is installed correctly and the optimized_kernels_is_available() test passes, then do the following:
Install makani but also pass --no-build-isolation to the pip install. Then see if it tries to reinstall either pytorch or torch harmonics. If not, it should be fine. Once you installed makani, repeat the above test. If that still passes, it should be able to use the optimized kernels.

Best regards
Thorsten