Skip to content

Commit

Permalink
Fix --cuaev-all-sms (aiqm#600)
Browse files Browse the repository at this point in the history
* Fix --cuaev-all-sms

* expecttest

* fix ci
  • Loading branch information
yueyericardo authored Oct 30, 2021
1 parent bf771af commit d1a6363
Show file tree
Hide file tree
Showing 3 changed files with 6 additions and 6 deletions.
2 changes: 1 addition & 1 deletion ci/install_dependencies.sh
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,6 @@

pip install --upgrade pip
pip install twine wheel
pip install --pre torch -f https://download.pytorch.org/whl/nightly/cu110/torch_nightly.html
pip install --pre torch -f https://download.pytorch.org/whl/nightly/cu112/torch_nightly.html
pip install -r test_requirements.txt
pip install -r docs_requirements.txt
9 changes: 4 additions & 5 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,10 +53,9 @@ def maybe_download_cub():
def cuda_extension(build_all=False):
import torch
from torch.utils.cpp_extension import CUDAExtension
SMs = None
SMs = []
print('-' * 75)
if not build_all:
SMs = []
devices = torch.cuda.device_count()
print('FAST_BUILD_CUAEV: ON')
print('This build will only support the following devices or the devices with same cuda capability: ')
Expand All @@ -71,13 +70,13 @@ def cuda_extension(build_all=False):
SMs.append(sm)

nvcc_args = ["-Xptxas=-v", '--expt-extended-lambda', '-use_fast_math']
if SMs:
if SMs and not ONLY_BUILD_SM80:
for sm in SMs:
nvcc_args.append(f"-gencode=arch=compute_{sm},code=sm_{sm}")
elif len(SMs) == 0 and ONLY_BUILD_SM80: # --cuaev --only-sm80
elif ONLY_BUILD_SM80: # --cuaev --only-sm80
nvcc_args.append("-gencode=arch=compute_80,code=sm_80")
else: # no gpu detected
print('NO gpu detected, will build for all SMs')
print('Will build for all SMs')
nvcc_args.append("-gencode=arch=compute_60,code=sm_60")
nvcc_args.append("-gencode=arch=compute_61,code=sm_61")
nvcc_args.append("-gencode=arch=compute_70,code=sm_70")
Expand Down
1 change: 1 addition & 0 deletions test_requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@ pillow
pkbar
pyyaml
pytest
expecttest

0 comments on commit d1a6363

Please sign in to comment.