-
Chris Klaiber authored
Summary: This unblocks building DRTK in a Docker build since BuildKit doesn't yet support GPUs: https://github.com/moby/buildkit/issues/1436 This works because CUDA is needed at run-time but not at build-time. When CUDA is present at build-time, the currently installed cards set which archs the extensions are built for. However, the archs are manually fixed in setup.py, so this commit also adds support for respecting TORCH_CUDA_ARCH_LIST which controls archs within torch.utils.cpp_extension.CUDAExtension Pull Request resolved: https://github.com/facebookresearch/DRTK/pull/5 Test Plan: Build without TORCH_CUDA_ARCH_LIST and observe using `ps` that compilation gets the flags specified in setup.py. Build with TORCH_CUDA_ARCH_LIST=Turing and TORCH_CUDA_ARCH_LIST=8.6 and observe using `ps` that compilation gets only the flags for the requested archs. Build within a Dockerfile using `docker-compose up --build`, which is using BuildKit, and observe that compilation succeeds and the library is usable when later run with CUDA GPUs available. Reviewed By: HapeMask Differential Revision: D63513797 Pulled By: podgorskiy fbshipit-source-id: 54ce6765ccc37317999f18690a73823eb074f08f
6ee72295
