mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
4df110b0a9
Summary: This is not actually needed and is causing a conda-forge confusion to do with python_abi - which needs users to have `-c conda-forge` when they install pytorch3d. Reviewed By: patricklabatut Differential Revision: D59587930 fbshipit-source-id: 961ae13a62e1b2b2ce6d8781db38bd97eca69e65
146 lines
4.6 KiB
Python
146 lines
4.6 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import os.path
|
|
import runpy
|
|
import subprocess
|
|
from typing import List
|
|
|
|
# required env vars:
|
|
# CU_VERSION: E.g. cu112
|
|
# JUST_TESTRUN: 1 to not set nvcc flags
|
|
# PYTORCH_VERSION: e.g. 1.12.0
|
|
# PYTHON_VERSION: e.g. 3.9
|
|
|
|
# should be run from pytorch3d root
|
|
|
|
CU_VERSION = os.environ["CU_VERSION"]
|
|
PYTORCH_VERSION = os.environ["PYTORCH_VERSION"]
|
|
pytorch_major_minor = tuple(int(i) for i in PYTORCH_VERSION.split(".")[:2])
|
|
source_root_dir = os.environ["PWD"]
|
|
|
|
|
|
def version_constraint(version):
|
|
"""
|
|
Given version "11.3" returns " >=11.3,<11.4"
|
|
"""
|
|
last_part = version.rindex(".") + 1
|
|
upper = version[:last_part] + str(1 + int(version[last_part:]))
|
|
return f" >={version},<{upper}"
|
|
|
|
|
|
def get_cuda_major_minor():
|
|
if CU_VERSION == "cpu":
|
|
raise ValueError("fn only for cuda builds")
|
|
if len(CU_VERSION) != 5 or CU_VERSION[:2] != "cu":
|
|
raise ValueError(f"Bad CU_VERSION {CU_VERSION}")
|
|
major = CU_VERSION[2:4]
|
|
minor = CU_VERSION[4]
|
|
return major, minor
|
|
|
|
|
|
def setup_cuda():
|
|
if CU_VERSION == "cpu":
|
|
return
|
|
major, minor = get_cuda_major_minor()
|
|
os.environ["CUDA_HOME"] = f"/usr/local/cuda-{major}.{minor}/"
|
|
os.environ["FORCE_CUDA"] = "1"
|
|
|
|
basic_nvcc_flags = (
|
|
"-gencode=arch=compute_50,code=sm_50 "
|
|
"-gencode=arch=compute_60,code=sm_60 "
|
|
"-gencode=arch=compute_70,code=sm_70 "
|
|
"-gencode=arch=compute_75,code=sm_75 "
|
|
"-gencode=arch=compute_50,code=compute_50"
|
|
)
|
|
if CU_VERSION == "cu102":
|
|
nvcc_flags = "-gencode=arch=compute_35,code=sm_35 " + basic_nvcc_flags
|
|
elif CU_VERSION < ("cu118"):
|
|
nvcc_flags = (
|
|
"-gencode=arch=compute_35,code=sm_35 "
|
|
+ "-gencode=arch=compute_80,code=sm_80 "
|
|
+ "-gencode=arch=compute_86,code=sm_86 "
|
|
+ basic_nvcc_flags
|
|
)
|
|
else:
|
|
nvcc_flags = (
|
|
"-gencode=arch=compute_80,code=sm_80 "
|
|
+ "-gencode=arch=compute_86,code=sm_86 "
|
|
+ "-gencode=arch=compute_90,code=sm_90 "
|
|
+ basic_nvcc_flags
|
|
)
|
|
|
|
if os.environ.get("JUST_TESTRUN", "0") != "1":
|
|
os.environ["NVCC_FLAGS"] = nvcc_flags
|
|
|
|
|
|
def setup_conda_pytorch_constraint() -> List[str]:
|
|
pytorch_constraint = f"- pytorch=={PYTORCH_VERSION}"
|
|
os.environ["CONDA_PYTORCH_CONSTRAINT"] = pytorch_constraint
|
|
if pytorch_major_minor < (2, 2):
|
|
os.environ["CONDA_PYTORCH_MKL_CONSTRAINT"] = "- mkl!=2024.1.0"
|
|
os.environ["SETUPTOOLS_CONSTRAINT"] = "- setuptools<70"
|
|
else:
|
|
os.environ["CONDA_PYTORCH_MKL_CONSTRAINT"] = ""
|
|
os.environ["SETUPTOOLS_CONSTRAINT"] = "- setuptools"
|
|
os.environ["CONDA_PYTORCH_BUILD_CONSTRAINT"] = pytorch_constraint
|
|
os.environ["PYTORCH_VERSION_NODOT"] = PYTORCH_VERSION.replace(".", "")
|
|
|
|
if pytorch_major_minor < (1, 13):
|
|
return ["-c", "pytorch"]
|
|
else:
|
|
return ["-c", "pytorch", "-c", "nvidia"]
|
|
|
|
|
|
def setup_conda_cudatoolkit_constraint():
|
|
if CU_VERSION == "cpu":
|
|
os.environ["CONDA_CPUONLY_FEATURE"] = "- cpuonly"
|
|
os.environ["CONDA_CUDATOOLKIT_CONSTRAINT"] = ""
|
|
return
|
|
os.environ["CONDA_CPUONLY_FEATURE"] = ""
|
|
|
|
if CU_VERSION in ("cu102", "cu110"):
|
|
os.environ["CONDA_CUB_CONSTRAINT"] = "- nvidiacub"
|
|
else:
|
|
os.environ["CONDA_CUB_CONSTRAINT"] = ""
|
|
|
|
major, minor = get_cuda_major_minor()
|
|
version_clause = version_constraint(f"{major}.{minor}")
|
|
if pytorch_major_minor < (1, 13):
|
|
toolkit = f"- cudatoolkit {version_clause}"
|
|
else:
|
|
toolkit = f"- pytorch-cuda {version_clause}"
|
|
os.environ["CONDA_CUDATOOLKIT_CONSTRAINT"] = toolkit
|
|
|
|
|
|
def do_build(start_args: List[str]):
|
|
args = start_args.copy()
|
|
|
|
test_flag = os.environ.get("TEST_FLAG")
|
|
if test_flag is not None:
|
|
args.append(test_flag)
|
|
|
|
args.extend(["-c", "bottler", "-c", "iopath", "-c", "conda-forge"])
|
|
args.append("--no-anaconda-upload")
|
|
args.extend(["--python", os.environ["PYTHON_VERSION"]])
|
|
args.append("packaging/pytorch3d")
|
|
print(args)
|
|
subprocess.check_call(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
args = ["conda", "build"]
|
|
setup_cuda()
|
|
|
|
init_path = source_root_dir + "/pytorch3d/__init__.py"
|
|
build_version = runpy.run_path(init_path)["__version__"]
|
|
os.environ["BUILD_VERSION"] = build_version
|
|
|
|
os.environ["SOURCE_ROOT_DIR"] = source_root_dir
|
|
args += setup_conda_pytorch_constraint()
|
|
setup_conda_cudatoolkit_constraint()
|
|
do_build(args)
|