Do not set ccbin in torch 1.7.x

Summary:
PyTorch versions 1.7.0 and 1.7.1 are between https://github.com/pytorch/pytorch/pull/43931 and https://github.com/pytorch/pytorch/pull/47404. In this gap, PyTorch always copies CC to nvcc_args, like PyTorch3D does. Newer nvcc versions are not happy with `-ccbin` being specified twice, even if it is specified twice the same. We update PyTorch3D so that it doesn't supply `-ccbin` in these cases.

Also tweak the detection of the current ccbin so that it is aware that `-ccbin foo` and `-ccbin=foo` are equivalent.

Reviewed By: theschnitz

Differential Revision: D25825468

fbshipit-source-id: b04e7718cf01820649518eedda99c399c732e8af
This commit is contained in:
Jeremy Reizenstein 2021-01-11 15:48:41 -08:00 committed by Facebook GitHub Bot
parent da7884cef1
commit b24d89a283

View File

@ -5,12 +5,31 @@ import glob
import os import os
import runpy import runpy
import warnings import warnings
from typing import List, Optional
import torch import torch
from setuptools import find_packages, setup from setuptools import find_packages, setup
from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension
def get_existing_ccbin(nvcc_args: List[str]) -> Optional[str]:
"""
Given a list of nvcc arguments, return the compiler if specified.
Note from CUDA doc: Single value options and list options must have
arguments, which must follow the name of the option itself by either
one of more spaces or an equals character.
"""
last_arg = None
for arg in reversed(nvcc_args):
if arg == "-ccbin":
return last_arg
if arg.startswith("-ccbin="):
return arg[7:]
last_arg = arg
return None
def get_extensions(): def get_extensions():
this_dir = os.path.dirname(os.path.abspath(__file__)) this_dir = os.path.dirname(os.path.abspath(__file__))
extensions_dir = os.path.join(this_dir, "pytorch3d", "csrc") extensions_dir = os.path.join(this_dir, "pytorch3d", "csrc")
@ -61,13 +80,18 @@ def get_extensions():
# This is needed for pytorch 1.6 and earlier. See e.g. # This is needed for pytorch 1.6 and earlier. See e.g.
# https://github.com/facebookresearch/pytorch3d/issues/436 # https://github.com/facebookresearch/pytorch3d/issues/436
CC = os.environ.get("CC", None) # It is harmless after https://github.com/pytorch/pytorch/pull/47404 .
if CC is not None: # But it can be problematic in torch 1.7.0 and 1.7.1
CC_arg = "-ccbin={}".format(CC) if torch.__version__[:4] != "1.7.":
if CC_arg not in nvcc_args: CC = os.environ.get("CC", None)
if any(arg.startswith("-ccbin") for arg in nvcc_args): if CC is not None:
raise ValueError("Inconsistent ccbins") existing_CC = get_existing_ccbin(nvcc_args)
nvcc_args.append(CC_arg) if existing_CC is None:
CC_arg = "-ccbin={}".format(CC)
nvcc_args.append(CC_arg)
elif existing_CC != CC:
msg = f"Inconsistent ccbins: {CC} and {existing_CC}"
raise ValueError(msg)
extra_compile_args["nvcc"] = nvcc_args extra_compile_args["nvcc"] = nvcc_args