@@ -18,31 +18,35 @@ export USE_STATIC_NCCL=1
18
18
export ATEN_STATIC_CUDA=1
19
19
export USE_CUDA_STATIC_LINK=1
20
20
21
+ # Keep an array of cmake variables to add to
22
+ if [[ -z " $CMAKE_ARGS " ]]; then
23
+ # These are passed to tools/build_pytorch_libs.sh::build()
24
+ CMAKE_ARGS=()
25
+ fi
26
+ if [[ -z " $EXTRA_CAFFE2_CMAKE_FLAGS " ]]; then
27
+ # These are passed to tools/build_pytorch_libs.sh::build_caffe2()
28
+ EXTRA_CAFFE2_CMAKE_FLAGS=()
29
+ fi
30
+
31
+ # Determine CUDA version and architectures to build for
21
32
CUDA_VERSION=$( nvcc --version| tail -n1| cut -f5 -d" " | cut -f1 -d" ," )
33
+ echo " CUDA $CUDA_VERSION Detected"
22
34
23
35
export TORCH_CUDA_ARCH_LIST=" 3.5;5.0+PTX"
24
36
if [[ $CUDA_VERSION == " 8.0" ]]; then
25
- echo " CUDA 8.0 Detected"
26
37
export TORCH_CUDA_ARCH_LIST=" $TORCH_CUDA_ARCH_LIST ;6.0;6.1"
27
38
elif [[ $CUDA_VERSION == " 9.0" ]]; then
28
- echo " CUDA $CUDA_VERSION Detected"
29
39
export TORCH_CUDA_ARCH_LIST=" $TORCH_CUDA_ARCH_LIST ;6.0;7.0"
30
40
elif [[ $CUDA_VERSION == " 9.2" ]]; then
31
- echo " CUDA $CUDA_VERSION Detected"
32
41
export TORCH_CUDA_ARCH_LIST=" $TORCH_CUDA_ARCH_LIST ;6.0;6.1;7.0"
33
- fi
34
- echo $TORCH_CUDA_ARCH_LIST
35
-
36
- if [[ $CUDA_VERSION == " 8.0" ]]; then
37
- WHEELHOUSE_DIR=" wheelhouse80"
38
- elif [[ $CUDA_VERSION == " 9.0" ]]; then
39
- WHEELHOUSE_DIR=" wheelhouse90"
40
- elif [[ $CUDA_VERSION == " 9.2" ]]; then
41
- WHEELHOUSE_DIR=" wheelhouse92"
42
+ # ATen tests can't build with CUDA 9.2 and the old compiler used here
43
+ EXTRA_CAFFE2_CMAKE_FLAGS+=(" -DATEN_NO_TEST=ON" )
42
44
else
43
45
echo " unknown cuda version $CUDA_VERSION "
44
46
exit 1
45
47
fi
48
+ echo $TORCH_CUDA_ARCH_LIST
49
+ WHEELHOUSE_DIR=" wheelhouse${CUDA_VERSION: 0: 1} {CUDA_VERSION:2:1}"
46
50
47
51
# Build for given Python versions, or for all in /opt/python if none given
48
52
if [[ -z " $DESIRED_PYTHON " ]]; then
@@ -90,7 +94,9 @@ for PYDIR in "${python_installations[@]}"; do
90
94
else
91
95
pip install numpy==1.11
92
96
fi
93
- time python setup.py bdist_wheel -d $WHEELHOUSE_DIR
97
+ time CMAKE_ARGS=${CMAKE_ARGS[@]} \
98
+ EXTRA_CAFFE2_CMAKE_FLAGS=${EXTRA_CAFFE2_CMAKE_FLAGS[@]} \
99
+ python setup.py bdist_wheel -d $WHEELHOUSE_DIR
94
100
done
95
101
96
102
popd
0 commit comments