Compiling pytorch with CUDA11.1+CUDNN 8.0.5, NO AVX with Ampere architecture

Hello everyone,

I am trying to compile pytorch on an Ubuntu 20 machine that does not have AVX support. The GPU on the server is based on the Ampere architecture with SM_86. I know that Cuda 11.1 and CUDNN 8.0.5 have SM_86 support and are not dependent on AVX (I successfully compiled Tensorflow, for example, and it works with the hardware).

I am looking for a combination of pytorch version and its dependencies that I compile with no AVX+the aforementioned CUDA/CUDNN versions. I cloned the github (recursively) and tried compiling different releases with the following options:

export TORCH_CUDA_ARCH_LIST=“8.6”
sudo cmake CC=/usr/bin/clang CXX=/usr/bin/clang++ -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc -DCMAKE_BUILD_TYPE=Debug -DPYTHON_EXECUTABLE=/usr/local/bin/python3 -DDISABLE_AVX2:BOOL=TRUE -DCXX_AVX2_FOUND:BOOL=FALSE -DC_AVX2_FOUND:BOOL=FALSE -DDISABLE_AVX512F:BOOL=TRUE -DDISABLE_AVX:BOOL=TRUE -DCXX_AVX_FOUND:BOOL=FALSE -DC_AVX_FOUND:BOOL=FALSE -DCMAKE_CXX_FLAGS:STRING=-std=c++14 …

I’ve tried both clang and gcc, but the compilation stops at different points with different release versions of pytorch. I get errors at pytorch-cpu or ATEN and sometimes one of the dependencies seems to assume that a later version of CUDA is present (judging by the header constants it is unable to find).

Does anyone have any suggestions?

Many thanks!